#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <string.h>
#include <stdlib.h>
#include <inttypes.h>
#include <unistd.h>
#include <stdio.h>
#include <errno.h>

const ssize_t SST_NEG1 = -1;
const size_t ST_NEG1 = -1;
const int SOCK_FALSE = 0;
const int SOCK_TRUE = 1;

#ifdef OFFICIAL
#define PORTNUM 12345
#define DEBUG 0
#else
#define PORTNUM 12345
#define DEBUG 1
#endif

typedef struct {
  uint8_t  version;
  uint8_t  ttl;
  uint16_t payload_length;
  uint32_t account_identifier;
  uint32_t source_address;
  uint32_t destination_address;
  uint16_t checksum;
  uint16_t protocol;
} cmsc417_message_header;

int main(int argc, char**argv) {
  int ret_val = 0;
  #ifdef OFFICIAL
  const int MAX_ARGS = 2;
  #else
  const int MAX_ARGS = 3;
  #endif
  /* 2^16 - 1 + 1, extra byte so we can definitely append null */
  const size_t MAX_SIZE = 65536;
  
  size_t msg_size = 0;
  void* msg_bytes = NULL;
  cmsc417_message_header head;
  struct sockaddr_in to_addr;
  memset(&to_addr, 0, sizeof(to_addr));
  int sock = 0;
  in_addr_t to_in_addr;
  memset(&to_in_addr, 0, sizeof(to_in_addr));
  socklen_t addr_size = 0;
  char* addr_str = NULL;
  ssize_t sent_size;
  int bind_ret = 0;
  void *buf = NULL;
  struct ip_mreq group;
  memset(&group, 0, sizeof(group));
  in_addr_t multi_in_addr;
  memset(&multi_in_addr, 0, sizeof(multi_in_addr));
  struct sockaddr_in multi_addr;
  memset(&multi_addr, 0, sizeof(multi_addr));
  size_t num_recv = 0;
  size_t tot_len = 0;
  
  if((2 <= argc) && (argc <= MAX_ARGS)) {
    /*************
     * Send
     *************/
    /* Declare and init message */
    head.version = 1;
    head.ttl = 1;
    /* Payload will be a string, argv[1], sans null  */
    head.payload_length = htons(strlen(argv[1]));
    head.account_identifier = htonl(6);
    head.source_address = htonl(0);
    head.destination_address = htonl(0);
    head.checksum = htons(0);
    head.protocol = htons(1);
    
    /* XXX this might bite us in the ass if the compiler pads the struct for alignment.
     * If it does, just calculate the unaligned size and use that for malloc and memcpy.
     */
    msg_size = sizeof(head) + strlen(argv[1]);
    msg_bytes = malloc(msg_size);
    if(NULL == msg_bytes) {
      fprintf(stderr, "Error allocating memory for message: %s\n", strerror(errno));
      goto END_SEND;
    }
    memset(msg_bytes, 0, sizeof(msg_size));
    memcpy(msg_bytes, &head, sizeof(head));
    memcpy(msg_bytes+sizeof(head), argv[0], strlen(argv[1]));
    
    if(argc >= 3) {
      addr_str = argv[2];
    } else {
      addr_str = "224.0.50.111";
    }
    to_in_addr = inet_addr(addr_str);
    if((INADDR_NONE == to_in_addr) && (0 != strncmp("255.255.255.255", addr_str, 15))) {
      fprintf(stderr, "Error converting address '%s': %s\n", addr_str, strerror(errno));
      goto END_SEND;
    }
    to_addr.sin_family = PF_INET;
    to_addr.sin_addr.s_addr = htonl(to_in_addr);
    to_addr.sin_port = htons(PORTNUM);

    /* SOCK_DGRAM and 0 might not be correct ... */
    sock = socket(PF_INET, SOCK_DGRAM, 0);
    if(-1 == sock) {
      fprintf(stderr, "Error creating send socket: %s\n", strerror(errno));
      goto END_SEND;
    }
    
    /* Note: saddr fields were initialized in network byte order */
    sent_size = sendto(sock, msg_bytes, msg_size, 0, (struct sockaddr *) &to_addr, sizeof(to_addr));
    if(DEBUG) { fprintf(stderr, "sent %zu out of %zu bytes\n", sent_size, msg_size); }
    if(SST_NEG1 == sent_size) {
      fprintf(stderr, "Error sending message: %s\n", strerror(errno));
    }
    close(sock);
    
    /* Jump Free Memory, Jump! */
  END_SEND:
    if(NULL != msg_bytes) { free(msg_bytes); msg_bytes = NULL; }
    
    /*************
     * Recieve
     *************/
    
    sock = socket(PF_INET, SOCK_DGRAM, 0);
    if(-1 == sock) {
      fprintf(stderr, "Error creating receive socket: %s\n", strerror(errno));
      ret_val = 10;
      goto END_RECV;
    }
    
    if(-1 == setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &SOCK_TRUE, sizeof(SOCK_TRUE))) {
      fprintf(stderr, "Error setting reuse addr on recv: %s\n", strerror(errno));
      close(sock);
      ret_val = 11;
      goto END_RECV;
    }
    /* This is needed on BSD (at least OS X) and shouldn't hurt */
    #ifdef SO_REUSEPORT
    if(-1 == setsockopt(sock, SOL_SOCKET, SO_REUSEPORT, &SOCK_TRUE, sizeof(SOCK_TRUE))) {
      fprintf(stderr, "Error setting reuse port on recv: %s\n", strerror(errno));
      ret_val = 12;
      close(sock);
      goto END_RECV;
    }
    #endif
    
    /* copy the address we sent to */
    multi_in_addr = inet_addr(addr_str);
    if((INADDR_NONE == multi_in_addr) && (0 != strncmp("255.255.255.255", addr_str, 15))) {
      fprintf(stderr, "Error converting address '%s': %s\n", addr_str, strerror(errno));
      goto END_RECV;
    }
    /* No htonl here. Weird? Yeah. */
    group.imr_multiaddr.s_addr = multi_in_addr;
    group.imr_interface.s_addr = htonl(INADDR_ANY);
    if(-1 == setsockopt(sock, IPPROTO_IP, IP_ADD_MEMBERSHIP, &group, sizeof(group))) {
      fprintf(stderr, "Error joining group on recv: %s\n", strerror(errno));
      close(sock);
      ret_val = 13;
      goto END_RECV;
    }

    multi_addr.sin_family = PF_INET;
    multi_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    multi_addr.sin_port = htons(PORTNUM);
    bind_ret = bind(sock, (struct sockaddr *) &multi_addr, sizeof(multi_addr));
    if(-1 == bind_ret) {
      fprintf(stderr, "Error binding to %s:%d: %s\n", addr_str, PORTNUM, strerror(errno));
      ret_val = 14;
      close(sock);
      goto END_RECV;
    }

    buf = malloc(MAX_SIZE);
    if(NULL == buf) {
      fprintf(stderr, "Error allocating memory for buffer: %s\n", strerror(errno));
      goto END_RECV;
    }
    while(1) {
      /* Safety first kids (debugging sanity, really) */
      memset(buf, 0, sizeof(buf));
      /* Why is the addr_size argument a non-const pointer?? */
      addr_size = sizeof(multi_addr);
      num_recv = recvfrom(sock, buf, sizeof(buf), MSG_WAITALL, (struct sockaddr *) &multi_addr, &addr_size);
      if(DEBUG) { fprintf(stderr, "recieved %zu bytes\n", num_recv); }
      if(-1 == num_recv) {
        fprintf(stderr, "Error receiving bytes: %s. Continuing ...\n", strerror(errno));
        continue;
      }
      if(num_recv >= sizeof(head)) {
        memcpy(&head, buf, sizeof(head));
        /* Validate */
        tot_len = ntohs(head.payload_length) + sizeof(head);
        if(DEBUG) {
          fprintf(stderr, "tot_len = %zu\n", tot_len);
          fprintf(stderr, "head.protocol = %d\n", ntohs(head.protocol));
          fprintf(stderr, "head.version = %d\n", head.version);
        }
        if((1 == ntohs(head.protocol)) && (1 == head.version) &&
           (tot_len == num_recv)) {
          ((char*)buf)[tot_len] = '\0'; /* ensure null termination */
          /* Is appending a newline okay? */
          printf("%s\n", (char*) buf+sizeof(head));
        } else {
          if(DEBUG) { fprintf(stderr, "Dropped packet\n"); }
        }
      }
    }
    close(sock);
  END_RECV:
    if(NULL != buf) { free(buf); buf = NULL; }
  } else {
    fprintf(stderr, "Usage:\n");
    #ifdef OFFICIAL
    fprintf(stderr, "\tone message\n");
    #else
    fprintf(stderr, "\tone message [address]\n");
    #endif
  }
  return ret_val;
}