#include "udp.h"

namespace udp {

int(*server_handler) (char* output_buffer, const char* input_buffer, int size) = NULL;
uint16_t server_port = 0;

void UDPPacket::setSrcPort(uint16_t port) {
  *(packet + 0) = port >> 8;
  *(packet + 1) = port & 0xff;
} 

void UDPPacket::setDstPort(uint16_t port) {
  *(packet + 2) = port >> 8;
  *(packet + 3) = port & 0xff;
}

uint16_t UDPPacket::getSrcPort() {
  return (*(packet + 0) << 8) + (uint8_t)*(packet + 1);
}
uint16_t UDPPacket::getDstPort() {
  return (*(packet + 2) << 8) + (uint8_t)*(packet + 3);
}

/* sets the payload and also updates the packet length field */
void UDPPacket::setPayload(const char* payload, int len) {
  memcpy(packet + 8, payload, len);
  *(packet + 4) = (len + 8) >> 8;
  *(packet + 5) = (len + 8) & 0xff;
}
const char* UDPPacket::getPayload() {
  return packet + 8;
}

/* computes the checksum and also updates it inside the UDP packet */
void UDPPacket::computeChecksum(const ipv6::IPv6Addr* src_addr, const ipv6::IPv6Addr* dst_addr, int len) {
  memset(packet + 6, 0, 2);
  uint16_t checksum = ipv6::compute_checksum(src_addr, dst_addr, ipv6::NH_UDP, len, packet, len);
  memcpy(packet + 6, &checksum, 2);
}

/* sets an handler (a function) which will be called every time there is a packet for a certain given port */
void set_handler(uint16_t port, int(* server_handler) (char* output_buffer, const char* input_buffer, int size)) {
  udp::server_port = port;
  udp::server_handler = server_handler;
}

/* handles an UDP packet, and (possibly) produces a reply UDP packet
 * returns length of reply packet, which is written in output_buffer, or -1 if there is no reply packet
*/
int handleUDP(char* output_buffer, const char* input_buffer, int ilen) {
  UDPPacket* incoming_packet = (UDPPacket*)input_buffer;
  UDPPacket* outgoing_packet = (UDPPacket*)output_buffer;
  int olen;

  if (incoming_packet->getDstPort() == server_port) { /* call user-defined server handler */
    if (server_port != 0 && server_handler != NULL) {
      olen = server_handler(outgoing_packet->getPayload(), incoming_packet->getPayload(), ilen - 8);
      
      outgoing_packet->setSrcPort(incoming_packet->getDstPort());
      outgoing_packet->setDstPort(incoming_packet->getSrcPort());
      outgoing_packet->setPayload(outgoing_packet->getPayload(), olen);
      outgoing_packet->computeChecksum(ipv6::packetin.getDstAddress(), ipv6::packetin.getSrcAddress(), olen + 8);
      return olen + 8;
    }
  } else {  /* default UDP application handler: echoes back the packet content */
    outgoing_packet->setSrcPort(incoming_packet->getDstPort());
    outgoing_packet->setDstPort(incoming_packet->getSrcPort());
    outgoing_packet->setPayload(incoming_packet->getPayload(), ilen - 8);
    outgoing_packet->computeChecksum(ipv6::packetin.getSrcAddress(), ipv6::packetin.getDstAddress(), ilen);
    return ilen;
  }
  
  return -1;
}

}