// Knows only protobuf_sgx objects, protobuf header.
// For socket programming
#include <sys/socket.h>
#include <stdlib.h>
#include <netinet/in.h>
#include <string.h>
#include <errno.h>
#include<unistd.h>
#include <stdio.h>
#include "dhmsgs.pb.h"
#include <google/protobuf/io/coded_stream.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
using namespace google::protobuf::io;
#include "protobufLAInitiator.h"

// TODO: Make these private functions
int read_protobuf_msg_from_fd(int accept_fd, google::protobuf::MessageLite& message)
{
  ZeroCopyInputStream* raw_input;
  CodedInputStream* coded_input;
  uint32_t size;
  CodedInputStream::Limit limit;
  raw_input = new FileInputStream(accept_fd);
  coded_input = new CodedInputStream(raw_input);
  if(!coded_input->ReadVarint32(&size))
  {
    printf("Error in reading size of msg");
    fflush(stdout);
    return -1;
  }
  //printf("size of msg was read to be %" PRIu32 " \n", size);
  fflush(stdout);
  limit = coded_input->PushLimit(size);
  if(!message.ParseFromCodedStream(coded_input))
  {
    printf("Error in parsing msg");
    fflush(stdout);
    return -1;
  }
  coded_input->PopLimit(limit);
  return 0;
}

// TODO: private functions
int write_protobuf_msg_to_fd(int accept_fd, google::protobuf::MessageLite& message)
{
  ZeroCopyOutputStream* raw_output = new FileOutputStream(accept_fd);
  CodedOutputStream* coded_output  = new CodedOutputStream(raw_output);
  coded_output->WriteVarint32(message.ByteSize());
  if(!message.SerializeToCodedStream(coded_output))
  {
    printf("SerializeToCodedStream failed");
    fflush(stdout);
    return -1;
  }
  // As per this - https://stackoverflow.com/questions/22881876/protocol-buffers-how-to-serialize-and-deserialize-multiple-messages-into-a-file?noredirect=1&lq=1
  // TODO: There may be a better way to do this - 1) this happens with every accept now and 2) make it happen on the stack vs heap - destructor will be called on return from this function (main) and the items will then be written out. (We probably don't want that, actually)
  delete coded_output;
  delete raw_output;
  fflush(stdout);
  return 0;
}

// Sets up a socket to bind and listen to the given port. Returns FD of the socket on success, -1 on failure (and prints a msg to stdout with the errno)
int set_up_socket(int port,   sockaddr_in* address)
{
  int server_fd = 0;

  // Creating socket file descriptor for listening for attestation requests.
  server_fd = socket(AF_INET, SOCK_STREAM | SOCK_CLOEXEC, 0);
  if (server_fd == -1)
  {
     printf("Error in creating a socket - %d", errno);
      return -1;
  }

  // Preparing the address struct for binding
  address->sin_family = AF_INET;
  address->sin_addr.s_addr = INADDR_ANY; // Todo: should this be localhost?
  address->sin_port = htons(port);
  //  memset(address->sin_zero,0,sizeof(address->sin_zero));

  socklen_t addrlen = sizeof(*address);
  // Binding
  if (bind(server_fd, (sockaddr*)address, addrlen)<0)
  {
      printf("Error in binding %d - port was %d - ", errno, port);
      return -1;
  }

  // Listening
  if (listen(server_fd, 128) < 0)
  {
      printf("Error in listening %d", errno);
      return -1;
  }

  return server_fd;
}

int local_attestation_initiator(int port, uint32_t own_enclave_id)
{
  // declare msg1, msg2, msg3 protobuf objects
  protobuf_sgx_dh_msg1_t protobuf_msg1;
  protobuf_sgx_dh_msg2_t protobuf_msg2;
  protobuf_sgx_dh_msg3_t protobuf_msg3;
  uint32_t protobuf_sgx_ret;

  // For socket to listen to the Apache enclave.
  int server_fd=0; int accept_fd = 0;
  struct sockaddr_in own_addr;
  struct sockaddr_storage apache_addr; socklen_t apache_addr_size = sizeof(apache_addr);
  uint32_t session_id; 
  //  int counter;
  server_fd=set_up_socket(port, &own_addr);
  if(server_fd==-1)
      return -1;
  printf("Successfully set up a socket to communicate with the Apache enclave.\n");
  fflush(stdout);

  protobuf_sgx_ret = generate_protobuf_dh_msg1(own_enclave_id, protobuf_msg1, &session_id);
  if(protobuf_sgx_ret != 0)
  {
    printf("Error in generate_protobuf_dh_msg1: 0x%x", protobuf_sgx_ret); fflush(stdout); return protobuf_sgx_ret;
  }


  accept_fd = accept(server_fd, (struct sockaddr *)&apache_addr,&apache_addr_size);
  if (accept_fd <0)
  {
        printf("Error in accepting %d", errno);
        return -1;
  }
  printf("Accepted fd\n"); fflush(stdout); 

  if(write_protobuf_msg_to_fd(accept_fd, protobuf_msg1)!=0)
    return -1;

  if(read_protobuf_msg_from_fd(accept_fd, protobuf_msg2)!=0)
    return -1;

  protobuf_sgx_ret = process_protobuf_dh_msg2_generate_protobuf_dh_msg3(own_enclave_id, protobuf_msg2, protobuf_msg3, &session_id);
  if(protobuf_sgx_ret != 0)
  {
      printf("Error in generate_protobuf_dh_msg2: 0x%x", protobuf_sgx_ret); fflush(stdout); return protobuf_sgx_ret;
  }

  if(write_protobuf_msg_to_fd(accept_fd, protobuf_msg3)!=0)
    return -1;

	printf("Successfully done Local attestation\n");
	fflush(stdout);
	return 0;
}