//
// Created by epagris on 2022.12.09..
//

#include <stdlib.h>
#include "dhcp.h"
#include "../../dynmem.h"
#include "../../utils.h"
#include "../../connection_block.h"
#include "../conn_blocks/udp_connblock.h"
#include "etherlib/global_state.h"

static struct {
    DhcpState state;
    void *buf;
    cbd desc;
    uint32_t tranId;
    EthInterface *intf;
} s;

static const uint8_t DHCP_MAGIC_COOKIE[] = {99, 130, 83, 99};

#define SNAME_LEN (64)
#define FILE_LEN (128)

// (Exact copy of the standard.)
typedef enum {
    DHCPDISCOVER = 1,
    DHCPOFFER = 2,
    DHCPREQUEST = 3,
    DHCPDECLINE = 4,
    DHCPACK = 5,
    DHCPNAK = 6,
    DHCPRELEASE = 7
} DhcpMsgType;

// DHCP option IDs
typedef enum {
    DHCP_OPT_SubnetMask = 0x1,
    DHCP_OPT_Router = 0x3,
    DHCP_OPT_DomainNameServer = 0x6,
    DHCP_OPT_RequestedIpAddress = 0x32,
    DHCP_OPT_IPAddrLeaseTime = 0x33,
    DHCP_OPT_MsgType = 0x35,
    DHCP_OPT_ServerId = 0x36,
    DHCP_OPT_ParamReqList = 0x37,
    DHCP_OPT_MaxMsgSize = 0x39,
    DHCP_OPT_ClientIdentifier = 0x3d,
    DHCP_OPT_End = 0xFF,
} DhcpOptionId;

#define MAX_DHCP_OPTION_LENGTH (16)

#define DHCP_HW_TYPE_ETHERNET (1)

// DHCP options
typedef struct DhcpOption_ {
    uint8_t id; ///< Options
    uint8_t length; ///< Option data length
    struct DhcpOption_ *next; /// next DHCP option
    uint8_t value[MAX_DHCP_OPTION_LENGTH]; ///< Option value
} DhcpOption;

static void dhcp_option_insert_msg_type(uint8_t **bufPtr, int msgType) {
    (*bufPtr)[0] = 0x35;
    (*bufPtr)[1] = 1;
    (*bufPtr)[2] = msgType;
    (*bufPtr) += 3;
}

static void dhcp_option_insert_max_msg_size(uint8_t **bufPtr, uint16_t maxSize) {
    (*bufPtr)[0] = 0x39;
    (*bufPtr)[1] = 2;
    (*bufPtr)[2] = (maxSize >> 8) & 0xFF;
    (*bufPtr)[3] = maxSize & 0xFF;
    (*bufPtr) += 4;
}

static void dhcp_option_insert_end(uint8_t **bufPtr) {
    (*bufPtr)[0] = 0xFF;
    (*bufPtr) += 1;
}

static void dhcp_read_next_option(const uint8_t **buf, DhcpOption *opt) {
    opt->id = (*buf)[0];
    if (opt->id != DHCP_OPT_End) {
        opt->length = (*buf)[1];
        memcpy(opt->value, (*buf) + 2, MIN(opt->length, MAX_DHCP_OPTION_LENGTH));
        (*buf) += 2 + opt->length;
    } else {
        (*buf) += 1;
    }
}

static void dhcp_insert_option(uint8_t **buf, const DhcpOption *opt) {
    (*buf)[0] = opt->id;
    if (opt->id != DHCP_OPT_End) {
        (*buf)[1] = opt->length;
        memcpy((*buf) + 2, opt->value, MIN(opt->length, MAX_DHCP_OPTION_LENGTH));
        (*buf) += 2 + opt->length;
    } else {
        (*buf) += 1;
    }
}

static void dhcp_free_opt_chain(DhcpOption *opt) {
    DhcpOption *iter = opt;
    while (iter) {
        DhcpOption *next = iter->next;
        dynmem_free(iter);
        iter = next;
    }
}

static const DhcpOption *dhcp_get_option(const DhcpOption *opts, DhcpOptionId id) {
    const DhcpOption *iter = opts;
    while (iter) {
        if (iter->id == id) {
            return iter;
        }
        iter = iter->next;
    }
    return NULL;
}

static void dhcp_send(const DhcpProps *props, const DhcpOption *opts) {
    // construct body
    uint8_t *buf = (uint8_t *) s.buf;
    memset(buf, 0, DHCP_MIN_PACKET_SIZE);
    FILL_BYTE_ADVANCE(buf, &(props->op));
    FILL_BYTE_ADVANCE(buf, &(props->htype));
    FILL_BYTE_ADVANCE(buf, &(props->hlen));
    FILL_BYTE_ADVANCE(buf, &(props->hops));
    FILL_DWORD_H2N_ADVANCE(buf, props->xid);
    FILL_WORD_H2N_ADVANCE(buf, props->secs);
    FILL_WORD_H2N_ADVANCE(buf, props->flags);
    FILL_ADVANCE(buf, &(props->ciaddr), 4);
    FILL_ADVANCE(buf, &(props->yiaddr), 4);
    FILL_ADVANCE(buf, &(props->siaddr), 4);
    FILL_ADVANCE(buf, &(props->giaddr), 4);
    FILL_ADVANCE(buf, props->chaddr, 16);
    buf += SNAME_LEN + FILE_LEN;
    FILL_ADVANCE(buf, DHCP_MAGIC_COOKIE, 4); // DHCP magic cookie

    // insert options
    const DhcpOption *iter = opts;
    while (iter) {
        dhcp_insert_option(&buf, iter);
        iter = iter->next;
    }

//    dhcp_option_insert_msg_type(&buf, DHCPDISCOVER);
//    dhcp_option_insert_max_msg_size(&buf, 1500);
//    dhcp_option_insert_end(&buf);

    // send packet
    udp_sendto(s.desc, s.buf, DHCP_MIN_PACKET_SIZE, IPv4_ANY_ADDR, DHCP_SERVER_PORT);
}

static void dhcp_parse(const uint8_t *buf, DhcpProps *props, DhcpOption **opts) {
    // parse body
    FETCH_BYTE_ADVANCE(&(props->op), buf);
    FETCH_BYTE_ADVANCE(&(props->htype), buf);
    FETCH_BYTE_ADVANCE(&(props->hlen), buf);
    FETCH_BYTE_ADVANCE(&(props->hops), buf);
    FETCH_DWORD_H2N_ADVANCE(&(props->xid), buf);
    FETCH_WORD_H2N_ADVANCE(&(props->secs), buf);
    FETCH_WORD_H2N_ADVANCE(&(props->flags), buf);
    FETCH_ADVANCE(&(props->ciaddr), buf, 4);
    FETCH_ADVANCE(&(props->yiaddr), buf, 4);
    FETCH_ADVANCE(&(props->siaddr), buf, 4);
    FETCH_ADVANCE(&(props->giaddr), buf, 4);
    FETCH_ADVANCE(props->chaddr, buf, 16);
    buf += SNAME_LEN + FILE_LEN;
    uint8_t magicCookie[sizeof(DHCP_MAGIC_COOKIE)];
    FETCH_ADVANCE(magicCookie, buf, 4); // DHCP magic cookie

    // parse options
    *opts = dynmem_alloc(sizeof(DhcpOption));
    (*opts)->next = NULL;
    dhcp_read_next_option(&buf, (*opts));
    DhcpOption *iter = *opts;

    while (iter->id != DHCP_OPT_End) {
        iter->next = dynmem_alloc(sizeof(DhcpOption));
        iter = iter->next;
        dhcp_read_next_option(&buf, iter);
        iter->next = NULL;
    }

}

#define UINT16_TO_BE_BYTES(u) ((u) >> 8) & 0xFF, ((u) & 0xFF),
#define UINT32_TO_BE_BYTES(u) ((u) >> 24) & 0xFF, ((u) >> 16) & 0xFF, ((u) >> 8) & 0xFF, ((u) & 0xFF),
#define IPv4_ADDR_TO_BE_BYTES(addr) ((addr) & 0xFF), (((addr) >> 8) & 0xFF), (((addr) >> 16) & 0xFF), (((addr) >> 24) & 0xFF),
#define HWADDR_TO_BE_BYTES(hwa) (hwa)[0], (hwa)[1], (hwa)[2], (hwa)[3], (hwa)[4], (hwa)[5],

static void dhcp_discover() {
    s.tranId = rand();
    DhcpProps props = {0};

    props.op = DHCP_BOOTREQUEST;
    props.htype = DHCP_HW_TYPE_ETHERNET;
    props.hlen = 6;
    props.hops = 0;
    props.xid = s.tranId;
    props.secs = 0;
    props.flags = 0;
    props.ciaddr = 0;
    props.yiaddr = 0;
    props.siaddr = 0;
    props.giaddr = 0;

    memcpy(props.chaddr, s.intf->mac, ETH_HW_ADDR_LEN);

    DhcpOption optEnd = {DHCP_OPT_End, 0, NULL};
    uint16_t maxSize = 1500;
    DhcpOption maxMsgSize = {DHCP_OPT_MaxMsgSize, 2, &optEnd, {UINT16_TO_BE_BYTES(maxSize)}};
    DhcpOption msgType = {DHCP_OPT_MsgType, 1, &maxMsgSize, {DHCPDISCOVER}};

    dhcp_send(&props, &msgType);
}

void dhcp_request(ip4_addr reqAddr, ip4_addr dhcpServerAddr) {
    DhcpProps props = {0};

    props.op = DHCP_BOOTREQUEST;
    props.htype = DHCP_HW_TYPE_ETHERNET;
    props.hlen = 6;
    props.hops = 0;
    props.xid = s.tranId;
    props.secs = 0;
    props.flags = 0;
    props.ciaddr = 0;
    props.yiaddr = 0;
    props.siaddr = 0;
    props.giaddr = 0;

    memcpy(props.chaddr, s.intf->mac, ETH_HW_ADDR_LEN);

    DhcpOption optEnd = {DHCP_OPT_End, 0, NULL};
    DhcpOption paramReq = {DHCP_OPT_ParamReqList, 4, &optEnd, {1, 3, 6, 51}}; // TODO...
    DhcpOption reqIp = {DHCP_OPT_RequestedIpAddress, 4, &paramReq, {IPv4_ADDR_TO_BE_BYTES(reqAddr)}};
    DhcpOption clId = {DHCP_OPT_ClientIdentifier, 7, &reqIp, {DHCP_HW_TYPE_ETHERNET, HWADDR_TO_BE_BYTES(s.intf->mac)}};
    DhcpOption msgType = {DHCP_OPT_MsgType, 1, &clId, {DHCPREQUEST}};

    dhcp_send(&props, &msgType);
}

static void dhcp_process(DhcpProps *props, DhcpOption *opts) {
    switch (s.state) {
        case DHCP_INIT:
            dhcp_discover(); // send discover message
            s.state = DHCP_SELECTING;
            break;
        case DHCP_SELECTING: {
            const DhcpOption *msgType = dhcp_get_option(opts, DHCP_OPT_MsgType);
            if (msgType->value[0] == DHCPOFFER) {
                ip4_addr addrOffer = props->yiaddr;
                ip4_addr serverAddr = props->siaddr;

                dhcp_request(addrOffer, serverAddr);

                s.state = DHCP_REQUESTING;
            }
        }
            break;
        case DHCP_REQUESTING: {
            const DhcpOption *opt = dhcp_get_option(opts, DHCP_OPT_MsgType);
            uint8_t msgType = opt->value[0];
            if (msgType == DHCPNAK) {
                //dhcp_discover();
                //s.state = DHCP_SELECTING;
            } else if (msgType == DHCPACK) {
                s.intf->ip = props->yiaddr; // fetch ip address

                opt = dhcp_get_option(opts, DHCP_OPT_Router); // get gateway/router address
                FETCH_DWORD(&(s.intf->router), opt->value);

                opt = dhcp_get_option(opts, DHCP_OPT_SubnetMask); // get subnet mask
                FETCH_DWORD(&(s.intf->netmask), opt->value);

                opt = dhcp_get_option(opts, DHCP_OPT_DomainNameServer); // get DNS
                FETCH_DWORD(&(s.intf->dns), opt->value);

                opt = dhcp_get_option(opts, DHCP_OPT_IPAddrLeaseTime); // fetch Lease Time
                uint32_t dhcpLeaseTime_s;
                FETCH_DWORD_H2N(&dhcpLeaseTime_s, opt->value);

                AlarmUserData params = { 0 };
                timer_sched_rel(E.tmr, ((int64_t)dhcpLeaseTime_s) * 1000000, NULL, params);

                MSG("DHCP done!\n");
                MSG("IP: ");
                PRINT_IPv4(s.intf->ip);
                MSG("\nRouter: ");
                PRINT_IPv4(s.intf->router);
                MSG("\nNetmask: ");
                PRINT_IPv4(s.intf->netmask);
                MSG("\nDNS: ");
                PRINT_IPv4(s.intf->dns);
                MSG("\nLease time: %u s\n", dhcpLeaseTime_s);
                MSG("\n");

                s.state = DHCP_BOUND;
            }
        }
            break;
        default:
            break;
    }
}

static int dhcp_resp_cb(const Pckt *pckt, PcktSieveLayerTag tag) {
    DhcpProps props;
    DhcpOption *opts = NULL;
    dhcp_parse(pckt->payload, &props, &opts);

    dhcp_process(&props, opts);

    dhcp_free_opt_chain(opts);

    return 0;
}

void dhcp_start() {
    s.state = DHCP_INIT;
    dhcp_process(NULL, NULL);
}

void dhcp_initiate(EthInterface *intf) {
    s.state = DHCP_INIT;
    s.buf = dynmem_alloc(DHCP_MIN_PACKET_SIZE);
    s.desc = udp_new_connblock(intf, IPv4_ANY_ADDR, DHCP_CLIENT_PORT, dhcp_resp_cb);
    s.intf = intf;
}