#include "DHCP_Server.h"

#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <cyw43_config.h>
#include <lwip/udp.h>
#include <lwip/ip_addr.h>
#include "DHCP_Config.h"

#define DHCPDISCOVER    (1)
#define DHCPOFFER       (2)
#define DHCPREQUEST     (3)
#define DHCPDECLINE     (4)
#define DHCPACK         (5)
#define DHCPNACK        (6)
#define DHCPRELEASE     (7)
#define DHCPINFORM      (8)

#define DHCP_OPT_PAD                (0)
#define DHCP_OPT_SUBNET_MASK        (1)
#define DHCP_OPT_ROUTER             (3)
#define DHCP_OPT_DNS                (6)
#define DHCP_OPT_HOST_NAME          (12)
#define DHCP_OPT_REQUESTED_IP       (50)
#define DHCP_OPT_IP_LEASE_TIME      (51)
#define DHCP_OPT_MSG_TYPE           (53)
#define DHCP_OPT_SERVER_ID          (54)
#define DHCP_OPT_PARAM_REQUEST_LIST (55)
#define DHCP_OPT_MAX_MSG_SIZE       (57)
#define DHCP_OPT_VENDOR_CLASS_ID    (60)
#define DHCP_OPT_CLIENT_ID          (61)
#define DHCP_OPT_END                (255)

#define MAC_LEN (6)

struct dhcp_server_lease_t {
    uint8_t mac[6];
    uint16_t expiry;
};

struct dhcp_server_t {
    ip_addr_t ip;
    ip_addr_t nm;
    struct dhcp_server_lease_t lease[DHCPS_MAX_IP];
    struct udp_pcb *pcb;
} dhcp_server;

struct dhcp_msg_t {
    uint8_t op; // message opcode
    uint8_t htype; // hardware address type
    uint8_t hlen; // hardware address length
    uint8_t hops;
    uint32_t xid; // transaction id, chosen by client
    uint16_t secs; // client seconds elapsed
    uint16_t flags;
    uint8_t ciaddr[4]; // client IP address
    uint8_t yiaddr[4]; // your IP address
    uint8_t siaddr[4]; // next server IP address
    uint8_t giaddr[4]; // relay agent IP address
    uint8_t chaddr[16]; // client hardware address
    uint8_t sname[64]; // server host name
    uint8_t file[128]; // boot file name
    uint8_t options[312]; // optional parameters, variable, starts with magic
};

#define DHCP_MSG_MIN_SIZE (240 + 3)

static void dhcp_socket_sendto(struct netif *nif, const void *buf, size_t len, uint32_t ip, uint16_t port)
{
    if(len > 0xffff)
    {
        len = 0xffff;
    }

    struct pbuf *p = pbuf_alloc(PBUF_TRANSPORT, len, PBUF_RAM);
    if(p == NULL)
    {
        puts("Pbuf allocation error");
        return;
    }

    memcpy(p->payload, buf, len);

    ip_addr_t dest;
    IP4_ADDR(ip_2_ip4(&dest), ip >> 24 & 0xff, ip >> 16 & 0xff, ip >> 8 & 0xff, ip & 0xff);

    err_t err;

    if(nif != NULL)
    {
        err = udp_sendto_if(dhcp_server.pcb, p, &dest, port, nif);
    }
    else
    {
        err = udp_sendto(dhcp_server.pcb, p, &dest, port);
    }

    pbuf_free(p);

    if(err != ERR_OK)
    {
        printf("Error sending DHCP packets : %d\n", err);
    }
}

static uint8_t *opt_find(uint8_t *opt, uint8_t cmd)
{
    for(int i = 0; i < 308 && opt[i] != DHCP_OPT_END; )
    {
        if (opt[i] == cmd)
        {
            return &opt[i];
        }

        i += 2 + opt[i + 1];
    }

    return NULL;
}

static void opt_write_n(uint8_t **opt, uint8_t cmd, size_t n, const void *data)
{
    uint8_t *o = *opt;
    *o++ = cmd;
    *o++ = n;
    memcpy(o, data, n);
    *opt = o + n;
}

static void opt_write_u8(uint8_t **opt, uint8_t cmd, uint8_t val)
{
    uint8_t *o = *opt;
    *o++ = cmd;
    *o++ = 1;
    *o++ = val;
    *opt = o;
}

static void opt_write_u32(uint8_t **opt, uint8_t cmd, uint32_t val)
{
    uint8_t *o = *opt;
    *o++ = cmd;
    *o++ = 4;
    *o++ = val >> 24;
    *o++ = val >> 16;
    *o++ = val >> 8;
    *o++ = val;
    *opt = o;
}

static void dhcp_server_process(void *arg __unused, struct udp_pcb *upcb __unused, struct pbuf *p, const ip_addr_t *src_addr __unused, u16_t src_port __unused)
{
    // This is around 548 bytes
    struct dhcp_msg_t dhcp_msg;

    if(p->tot_len < DHCP_MSG_MIN_SIZE)
    {
        goto ignore_request;
    }

    size_t len = pbuf_copy_partial(p, &dhcp_msg, sizeof(dhcp_msg), 0);
    if(len < DHCP_MSG_MIN_SIZE)
    {
        goto ignore_request;
    }

    dhcp_msg.op = DHCPOFFER;
    memcpy(&dhcp_msg.yiaddr, &ip4_addr_get_u32(ip_2_ip4(&dhcp_server.ip)), 4);

    uint8_t *opt = (uint8_t *)&dhcp_msg.options;
    opt += 4; // assume magic cookie: 99, 130, 83, 99

    uint8_t *msgtype = opt_find(opt, DHCP_OPT_MSG_TYPE);
    if(msgtype == NULL)
    {
        // A DHCP package without MSG_TYPE?
        goto ignore_request;
    }

    switch(msgtype[2])
    {
      case DHCPDISCOVER:
      {
        int yi = DHCPS_MAX_IP;

        for(int i = 0; i < DHCPS_MAX_IP; ++i)
        {
            if(memcmp(dhcp_server.lease[i].mac, dhcp_msg.chaddr, MAC_LEN) == 0)
            {
                // MAC match, use this IP address
                yi = i;
                break;
            }

            if(yi == DHCPS_MAX_IP)
            {
                // Look for a free IP address
                if (memcmp(dhcp_server.lease[i].mac, "\x00\x00\x00\x00\x00\x00", MAC_LEN) == 0)
                {
                    // IP available
                    yi = i;
                }

                uint32_t expiry = dhcp_server.lease[i].expiry << 16 | 0xffff;
                if((int32_t)(expiry - cyw43_hal_ticks_ms()) < 0)
                {
                    // IP expired, reuse it
                    memset(dhcp_server.lease[i].mac, 0, MAC_LEN);
                    yi = i;
                }
            }
        }

        if(yi == DHCPS_MAX_IP)
        {
            // No more IP addresses left
            goto ignore_request;
        }

        dhcp_msg.yiaddr[3] = DHCPS_BASE_IP + yi;
        opt_write_u8(&opt, DHCP_OPT_MSG_TYPE, DHCPOFFER);

        break;
      }

      case DHCPREQUEST:
      {
        uint8_t *o = opt_find(opt, DHCP_OPT_REQUESTED_IP);
        if(o == NULL)
        {
            // Should be NACK
            goto ignore_request;
        }
        if(memcmp(o + 2, &ip4_addr_get_u32(ip_2_ip4(&dhcp_server.ip)), 3) != 0)
        {
            // Should be NACK
            goto ignore_request;
        }

        uint8_t yi = o[5] - DHCPS_BASE_IP;

        if(yi >= DHCPS_MAX_IP)
        {
            // Should be NACK
            goto ignore_request;
        }

        if(memcmp(dhcp_server.lease[yi].mac, dhcp_msg.chaddr, MAC_LEN) == 0)
        {
            // MAC match, ok to use this IP address
        }
        else if(memcmp(dhcp_server.lease[yi].mac, "\x00\x00\x00\x00\x00\x00", MAC_LEN) == 0)
        {
            // IP unused, ok to use this IP address
            memcpy(dhcp_server.lease[yi].mac, dhcp_msg.chaddr, MAC_LEN);
        }
        else
        {
            // IP already in use
            // Should be NACK
            goto ignore_request;
        }

        dhcp_server.lease[yi].expiry = (cyw43_hal_ticks_ms() + DEFAULT_LEASE_TIME_S * 1000) >> 16;
        dhcp_msg.yiaddr[3] = DHCPS_BASE_IP + yi;

        opt_write_u8(&opt, DHCP_OPT_MSG_TYPE, DHCPACK);

        printf("DHCP Server: client connected: MAC=%02x:%02x:%02x:%02x:%02x:%02x IP=%u.%u.%u.%u\n",
            dhcp_msg.chaddr[0], dhcp_msg.chaddr[1], dhcp_msg.chaddr[2], dhcp_msg.chaddr[3], dhcp_msg.chaddr[4], dhcp_msg.chaddr[5],
            dhcp_msg.yiaddr[0], dhcp_msg.yiaddr[1], dhcp_msg.yiaddr[2], dhcp_msg.yiaddr[3]);

        break;
    }

      default:
        goto ignore_request;
    }

    opt_write_n(&opt, DHCP_OPT_SERVER_ID, 4, &ip4_addr_get_u32(ip_2_ip4(&dhcp_server.ip)));
    opt_write_n(&opt, DHCP_OPT_SUBNET_MASK, 4, &ip4_addr_get_u32(ip_2_ip4(&dhcp_server.nm)));
    opt_write_n(&opt, DHCP_OPT_ROUTER, 4, &ip4_addr_get_u32(ip_2_ip4(&dhcp_server.ip))); // aka gateway; can have multiple addresses
    opt_write_n(&opt, DHCP_OPT_DNS, 4, &ip4_addr_get_u32(ip_2_ip4(&dhcp_server.ip))); // this server is the dns
    opt_write_u32(&opt, DHCP_OPT_IP_LEASE_TIME, DEFAULT_LEASE_TIME_S);
    *opt++ = DHCP_OPT_END;
    struct netif *nif = ip_current_input_netif();
    dhcp_socket_sendto(nif, &dhcp_msg, opt - (uint8_t *)&dhcp_msg, 0xffffffff, PORT_DHCP_CLIENT);

ignore_request:
    pbuf_free(p);
}

int dhcp_server_init(void)
{
    dhcp_server.ip.addr = DHCP_SERVER_IP;
    dhcp_server.nm.addr = DHCP_SERVER_MASK;
    memset(dhcp_server.lease, 0, sizeof(dhcp_server.lease));

    dhcp_server.pcb = udp_new();
    if(dhcp_server.pcb == NULL)
    {
        puts("Error creating DHCP server");
        return PICO_ERROR_GENERIC;
    }

    // Register callback
    udp_recv(dhcp_server.pcb, dhcp_server_process, NULL);

    err_t err = udp_bind(dhcp_server.pcb, IP_ANY_TYPE, PORT_DHCP_SERVER);
    if(err != ERR_OK)
    {
        printf("Error binding DHCP server : %d\n", err);
        return err;
    }

    return PICO_OK;
}

void dhcp_server_exit(void)
{
    if(dhcp_server.pcb != NULL)
    {
        udp_remove(dhcp_server.pcb);
        dhcp_server.pcb = NULL;
    }
}