#include <stdint.h>
#include <string.h>

#include <networkAbstractionLayer_impl.h>
#include <logger.h>
#include <PontCoopScheduler.h>
#include <wizHelper.h>
#include <wizchip_conf.h>
#include <socket.h>
#include <wizchip_conf.h>
#include <config.h>

static t_configBlock *config;


const uint16_t MAX_RECV_RETRY_COUNT = 100;
const uint64_t UNIX_NTP_EPOCH_DIFF = 2208988800;

const uint8_t SEND_LI_VN_MODE = 0xe3; // LI: unknown (3), VN: 4, Mode: Client (3)
typedef struct {
    uint8_t   li_vn_mode;
	uint8_t   stratum;
	uint8_t   poll;
	uint8_t   precision;
	uint32_t  rootdelay;
	uint32_t  rootdisp;
	uint32_t  refid;
	uint64_t  reftime;
	uint64_t  org;
	uint64_t  rec;
	uint64_t  xmt;
} ntpMsg_t;

extern uint8_t SNTP_SOCK;

const uint16_t NTP_PORT = 123;

enum { 
    SNTP_STATE_IDLE, 
    SNTP_STATE_START,
    SNTP_STATE_SEND,
    SNTP_STATE_RECV,
    SNTP_STATE_ERROR,
    SNTP_STATE_DONE 
} sntpState = SNTP_STATE_IDLE;
uint64_t seconds;

static void networkSntpEngine(void *handle) {
    static uint16_t retryCount = 0;

    coloredMsg(LOG_BLUE, "nes, %u", sntpState);

    switch (sntpState) {
        case SNTP_STATE_START:
            coloredMsg(LOG_BLUE, "nes, about to send");
            uint8_t ntpAddr[4];
            if (! wizDnsQuery(config->ntpServer, ntpAddr)) {
                coloredMsg(LOG_BLUE, "nes, failed to resolve ntp server");
                sntpState = SNTP_STATE_ERROR;
            } else {
                socket(SNTP_SOCK, Sn_MR_UDP, NTP_PORT, 0);
                retryCount = 0;
                uint8_t sockState = getSn_SR(SNTP_SOCK);
                if (sockState == SOCK_UDP) {
                    ntpMsg_t ntpMsg;
                    memset(&ntpMsg, 0, sizeof(ntpMsg));
                    ntpMsg.li_vn_mode = SEND_LI_VN_MODE;
                    sendto(SNTP_SOCK, (uint8_t*)&ntpMsg, sizeof(ntpMsg), ntpAddr, NTP_PORT);
                    coloredMsg(LOG_BLUE, "nes, sent");
                    sntpState = SNTP_STATE_RECV;
                    schAdd(networkSntpEngine, NULL, 1000, 0);
                } else {
                    coloredMsg(LOG_BLUE, "nes, socket in unexpected state: %d", sockState);
                    sntpState = SNTP_STATE_ERROR;
                }
            }
            break;
        case SNTP_STATE_RECV:
            coloredMsg(LOG_BLUE, "nes, check receive");
            uint16_t recvLen = getSn_RX_RSR(SNTP_SOCK);
            if (recvLen == 0) {
                retryCount += 1;
                if (retryCount > MAX_RECV_RETRY_COUNT) {
                    coloredMsg(LOG_BLUE, "nes max retry count reached, failed");
                    sntpState = SNTP_STATE_ERROR;
                } else {
                    coloredMsg(LOG_BLUE, "nes, nothing received yet, try again");
                    schAdd(networkSntpEngine, NULL, 100, 0);
                }
            } else if (recvLen == sizeof(ntpMsg_t)) {
                ntpMsg_t ntpMsg;
                memset(&ntpMsg, 0, sizeof(ntpMsg_t));
                uint8_t srcAddr[4];
                uint16_t srcPort;
                recvfrom(SNTP_SOCK, (uint8_t*)&ntpMsg, sizeof(ntpMsg_t), srcAddr, &srcPort);
                close(SNTP_SOCK);
                coloredMsg(LOG_BLUE, "nes, msg received from %d.%d.%d.%d:%d", 
                           srcAddr[0], srcAddr[1], srcAddr[2], srcAddr[3],
                           srcPort);
                coloredMsg(LOG_BLUE, "nes, received in the %d. cycles", retryCount);
                seconds = ntpMsg.rec - UNIX_NTP_EPOCH_DIFF;
                coloredMsg(LOG_BLUE, "nes, seconds: %lu", seconds);
                sntpState = SNTP_STATE_DONE;
            } else {
                coloredMsg(LOG_BLUE, "nes, invalid number of octets received: %d", recvLen);
                sntpState = SNTP_STATE_ERROR;
            }
            break;
        default:
            coloredMsg(LOG_BLUE, "nes, unexpected state");
    }
}

uint64_t networkSntpQuery() {
    uint64_t res = 0;
    if (sntpState == SNTP_STATE_IDLE) {
        sntpState = SNTP_STATE_START;
        schAdd(networkSntpEngine, NULL, 1, 0);
    } else if (sntpState == SNTP_STATE_ERROR) {
        sntpState = SNTP_STATE_IDLE;
    } else if (sntpState == SNTP_STATE_DONE) {
        sntpState = SNTP_STATE_IDLE;
        res = seconds;
    }
    return res;
}


extern uint8_t SINK_SOCK;

int8_t networkUdpSend(char *hostname, uint16_t port, uint8_t *buf, uint16_t bufLen) {
    uint8_t sinkAddr[4];
    if (! wizDnsQuery(hostname, sinkAddr)) {
        coloredMsg(LOG_BLUE, "nus, failed to resolve sink server name");
        return -1;
    } else {
        coloredMsg(LOG_BLUE, "nus, sink server at %d.%d.%d.%d", sinkAddr[0], sinkAddr[1], sinkAddr[2], sinkAddr[3]);
    }

    socket(SINK_SOCK, Sn_MR_UDP, port, 0);
    uint8_t sockState = getSn_SR(SINK_SOCK);
    if (sockState == SOCK_UDP) {
        sendto(SINK_SOCK, buf, bufLen, sinkAddr, port);
        coloredMsg(LOG_BLUE, "nus, sent");
    } else {
        coloredMsg(LOG_BLUE, "nus, socket in unexpected state: %d", sockState);
        return -2;
    }

    close(SINK_SOCK);

    return 1;
}

void networkImplInit() {
    config = getConfig();
    wizInit();
}