#include "contiki.h"
#include "ecc.h"
#include "ecc_dh.h"
#include "ecc_dsa.h"
#include "test_ecc_utils.h"
#include "test_uti.h"
#include "test_ecc_utils.h"
#include "test_uti.h"
#include "myencrypt.h"
#include "sha256.h"
#include "sys/energest.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

#define SEND_INTERVAL		  (60 * CLOCK_SECOND)

static struct simple_udp_connection udp_conn;

PROCESS(Secure_Client, "Secure Client Service");
AUTOSTART_PROCESSES(&Secure_Client);

//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");

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

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

	#if LLSEC802154_CONF_ENABLED
	  LOG_INFO_(" LLSEC LV:%d", uipbuf_get_attr(UIPBUF_ATTR_LLSEC_LEVEL));
	#endif
	  LOG_INFO_("\n");
}

PROCESS_THREAD(Secure_Client, ev, data)
{
    static struct etimer periodic_timer;
  	static unsigned count;
  	uip_ipaddr_t dest_ipaddr;

  	PROCESS_BEGIN();
    struct GenKeyPair keypair1;

    //uint8_t ciphertext[32];
    //uint8_t plaintext[32];
    //uint8_t key[NUM_ECC_BYTES];
    //uint8_t signtext[2 * NUM_ECC_BYTES];

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

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

    keypair1  = ECC_GenKeyPair();

  	etimer_set(&periodic_timer, random_rand() % SEND_INTERVAL);
  	while(1) 
  	{
    	PROCESS_WAIT_EVENT_UNTIL(etimer_expired(&periodic_timer));

    	if(NETSTACK_ROUTING.node_is_reachable() && NETSTACK_ROUTING.get_root_ipaddr(&dest_ipaddr)) {
      	/* Send to DAG root */
      	LOG_INFO("Sending request %u to ", count);
      	LOG_INFO_6ADDR(&dest_ipaddr);
      	LOG_INFO_("\n");

      	simple_udp_sendto(&udp_conn, keypair1.pubkey, sizeof(keypair1.pubkey), &dest_ipaddr);
      	count++;
    } 
    
    else 
    {
      LOG_INFO("Not reachable yet\n");
    }

    /* Add some jitter */
    etimer_set(&periodic_timer, SEND_INTERVAL
      - CLOCK_SECOND + (random_rand() % (2 * CLOCK_SECOND)));
  }

  PROCESS_END();
}