#include "contiki.h"
#include "ecc.h"
#include "ecc_dh.h"
#include "test_ecc_utils.h"
#include "test_uti.h"

#include "net/routing/routing.h"
#include "random.h"
#include "net/netstack.h"
#include "net/ipv6/simple-udp.h"

#include "sys/log.h"
#define LOG_MODULE "App"
#define LOG_LEVEL LOG_LEVEL_INFO

#include <assert.h>
#include <stdio.h> 
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>


#define WITH_SERVER_REPLY  1
#define UDP_CLIENT_PORT	8765
#define UDP_SERVER_PORT	5678

static struct simple_udp_connection udp_conn;

PROCESS(Secure_Server, "Secure Server Service");
AUTOSTART_PROCESSES(&Secure_Server);

static unsigned long
to_seconds(uint64_t time)
{
  return (unsigned long)(time/ENERGEST_SECOND);
}

struct GenKeyPair
{
  uint8_t privkey[NUM_ECC_BYTES];
  uint8_t pubkey[2 * NUM_ECC_BYTES];
};

struct GenKeyPair ECC_GenKeyPair()
{
  static struct GenKeyPair genpair1;

  const struct uECC_Curve_t * curve = uECC_secp256r1();

  clock_time_t start_time = clock_time();
  uECC_make_key(genpair1.pubkey, genpair1.privkey, curve);
  clock_time_t end_time = clock_time();

  unsigned long time_taken = end_time - start_time;
  printf("Time Taken to Generate Key Pair: %lu ticks\n", time_taken);
  
  show_str("ECC Private Key:", genpair1.privkey, sizeof(genpair1.privkey));
    show_str("ECC Public Key:", genpair1.pubkey, sizeof(genpair1.pubkey));
    printf("\n");


  return genpair1;
}

int ECDH_GenSharedKey(uint8_t *privkey, uint8_t *pubkey, uint8_t *secret, int secretlen)
{
  const struct uECC_Curve_t * curve = uECC_secp256r1();

  clock_time_t start_time = clock_time();
  uECC_shared_secret(pubkey, privkey, secret, curve);
  clock_time_t end_time = clock_time();

  unsigned long time_taken = end_time - start_time;
  printf("Time Taken to Generate Shared Secret: %lu ticks\n", time_taken);

  show_str("ECDH Shared Key: ", secret, secretlen);
  printf("\n\n");

  return 0;
}

int AES_Encrypt(uint8_t *m, uint8_t * ciphertext, TCAesKeySched_t s, int len)
{
  clock_time_t start_time = clock_time();
  tc_aes_encrypt(ciphertext, m, s);
  
  clock_time_t end_time = clock_time();
  
  unsigned long time_taken = end_time - start_time;
  printf("Time Taken to Encrypt Message: %lu ticks\n", time_taken);

  printf("Start: %lu ticks\n", start_time);
  printf("End: %lu ticks\n", end_time);

  show_str("Encrypted Message: ", ciphertext, len);
  printf("\n");

  return 0;
}

int AES_Decrypt(uint8_t * ciphertext, uint8_t * plaintext, TCAesKeySched_t s, int len)
{
  clock_time_t start_time = clock_time();
  tc_aes_decrypt(plaintext, ciphertext, s);
  
  clock_time_t end_time = clock_time();
  
  unsigned long time_taken = end_time - start_time;
  printf("Time Taken to Decrypt Message: %lu ticks\n", time_taken);

  printf("Start: %lu ticks\n", start_time);
  printf("End: %lu ticks\n", end_time);

  show_str("Decrypted Message: ", plaintext, len);
  printf("\n");

  return 0;
}

int ECDSA_sign(uint8_t * privkey, char * msg, uint8_t * signtext, int signlen)
{
  uint8_t hashtag[32];
  const struct uECC_Curve_t * curve = uECC_secp256r1();
  struct tc_sha256_state_struct s;
  (void)tc_sha256_init(&s);
  tc_sha256_update(&s, (const uint8_t *) msg, strlen(msg));
    (void)tc_sha256_final(hashtag, &s);
  

  clock_time_t start_time = clock_time();
  uECC_sign(privkey, hashtag, sizeof(hashtag), signtext, curve);
  clock_time_t end_time = clock_time();
  
  unsigned long time_taken = end_time - start_time;
  printf("Time Taken to Sign Message: %lu ticks\n", time_taken);

  show_str("Signed Message: ", signtext, signlen);
  printf("\n");

  return 0;
}

int ECDSA_verify(uint8_t * pubkey, char * msg, uint8_t * signtext)
{
  uint8_t hashtag[32];
  const struct uECC_Curve_t * curve = uECC_secp256r1();
  struct tc_sha256_state_struct s;
  (void)tc_sha256_init(&s);
  tc_sha256_update(&s, (const uint8_t *) msg, strlen(msg));
    (void)tc_sha256_final(hashtag, &s);

  clock_time_t start_time = clock_time();
  if(!uECC_verify(pubkey, hashtag, sizeof(hashtag), signtext, curve))
  {
    printf("Verification Unsuccessful! Check Signature!\n");
  }
  else
  {
    
    printf("Successful Verification! Signature is genuine\n");
  }
  clock_time_t end_time = clock_time(); 
  unsigned long time_taken = end_time - start_time;
  printf("Time Taken to Verify Signature: %lu ticks\n", time_taken);

  return 0;
}


static void
udp_rx_callback(struct simple_udp_connection *c,
	         const uip_ipaddr_t *sender_addr,
	         uint16_t sender_port,
	         const uip_ipaddr_t *receiver_addr,
	         uint16_t receiver_port,
	         const uint8_t *data,
	         uint16_t datalen)
{
 	LOG_INFO("Received request from ");
  LOG_INFO_6ADDR(sender_addr);
  LOG_INFO_("\n");
  uint8_t Key[2 * NUM_ECC_BYTES];

  memcpy(Key, (uint8_t *)data, sizeof(Key));

  const struct uECC_Curve_t * curve = uECC_secp256r1();

  show_str("Public Key Received", Key, sizeof(Key));
  printf("\n");

  uECC_shared_secret(Key, privkey, secret, curve);
  show_str("Shared Secret", secret, sizeof(secret));
  printf("\n");


	#if WITH_SERVER_REPLY
  	LOG_INFO("Exchange Public Key with Node 1 ...............................\n");
  	printf("\n");

  	simple_udp_sendto(&udp_conn, (const void *)pubkey, datalen, sender_addr);

  #endif /* WITH_SERVER_REPLY */
}

PROCESS_THREAD(Secure_Server, ev, data)
{	

	PROCESS_BEGIN();

  srand(5728);

  /* Setup of the Cryptographically Secure PRNG. */
  uECC_set_rng(&default_CSPRNG);

  const struct uECC_Curve_t * curve = uECC_secp256r1();

  uECC_make_key(pubkey, privkey, curve);

  show_str("Private Key", privkey, sizeof(privkey));
  printf("\n");
  show_str("Public Key", pubkey, sizeof(pubkey));
  printf("\n");

	/* Initialize DAG root */
	NETSTACK_ROUTING.root_start();

	/* Initialize UDP connection */
	simple_udp_register(&udp_conn, UDP_SERVER_PORT, NULL, UDP_CLIENT_PORT, udp_rx_callback);

	PROCESS_END();
}