#include "udp.h"

namespace udp {

void(*server_handler) (const char* buffer, int size) = NULL;
uint16_t server_port = 0;

UDPPacket::UDPPacket() {
  memset(packet, 0, LEN);
}

UDPPacket::UDPPacket(const char* buffer, int len) {
  memcpy(packet, buffer, len);
}

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);
}

void UDPPacket::setPayload(const char* payload, int len) {
  memcpy(packet + 8, payload, len);
  *(packet + 4) = (len + 8) >> 8;
  *(packet + 5) = (len + 8) & 0xff;
}

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);
}

void set_handler(uint16_t port, void(* server_handler) (const char* buffer, int size)) {
  udp::server_port = port;
  udp::server_handler = server_handler;
}

void handleUDP(const ipv6::IPv6Addr* src_addr, const ipv6::IPv6Addr* dst_addr, const char* flow, char* packet, int len) {
  UDPPacket udp_packet(packet, len);

  if (udp_packet.getDstPort() == server_port) { /* call user-defined server handler */
    if (server_port != 0 && server_handler != NULL) {
      server_handler(packet + 8, len - 8);
    }
  } else {  /* echo default UDP application handler */
    /* swap ports */
    uint16_t tmp;
    tmp = udp_packet.getSrcPort();
    udp_packet.setSrcPort(udp_packet.getDstPort());
    udp_packet.setDstPort(tmp);
  
    udp_packet.computeChecksum(src_addr, dst_addr, len);
  
    /* prepare IPv6 packet */
    ipv6::IPv6Packet reply_packet;
    reply_packet.setPayload((char*)&udp_packet, len);
    reply_packet.setFlow(flow);
    
    reply_packet.setSrcAddress(dst_addr);
    reply_packet.setDstAddress(src_addr);
    reply_packet.setNextHeader(ipv6::NH_UDP);
    slip::send((void*)&reply_packet, 40 + len);
  }
}

}