#include "tcp_connblock.h"

#include <stddef.h>
#include <stdlib.h>

#include "../packet_parsers/packet_parsers.h"
#include "../../utils.h"
#include "../../dynmem.h"
#include "../../pckt_assembler.h"
#include "../../eth_interface.h"
#include "ipv4_connblock.h"
#include "../../gen_queue.h"
#include "etherlib_options.h"
#include "../../prefab/conn_blocks/tcp/tcp_window.h"
#include "../../global_state.h"

static char *TCP_STATE_NAMES[] = {
        "CLOSED", "LISTEN", "SYN_RCVD", "SYN_SENT", "ESTAB", "FIN_WAIT_1",
        "FIN_WAIT_2", "CLOSE_WAIT", "CLOSING", "LAST_ACK", "TIME_WAIT"
};

int tcp_send_segment(const struct ConnBlock_ *connBlock, TcpFlag flags, TcpOption *opts, const uint8_t *data, uint32_t size);

static bool filtTcp(const PcktSieveFilterCondition *filtCond, const PcktProps *contProps, const PcktProps *ownProps, EthInterface *intf) {
    IPv4Props *ipProps = (IPv4Props *) contProps;
    TcpProps *tcpProps = (TcpProps *) ownProps;

    return ipProps->Protocol == ETH_TCP_PACKET_CLASS && (TCP_PORT_FROM_FILTCOND(filtCond) == tcpProps->DestinationPort);
}

/**
 * TCP state.
 */
typedef struct {
    uint16_t localPort;
    uint16_t remotePort;
    ip4_addr remoteAddr;
    TcpConnectionState connState;
    uint16_t localMSS;
    uint16_t remoteMSS;
    uint32_t sequenceNumber; // seq. num of next transmitted octet
    uint32_t ackNumber; // ack. number sent to the remote side (seq. num of next expected octet)
    uint16_t window;
    TcpWindow *txWin;
    ConnBlock connBlock;
    cbd d; // connection block descriptor
    uint8_t retryLimit;
    uint8_t retries;
    uint8_t retryTO; // retry timeout
    bool8_t debug;
    SieveCallBackFn userCb;
} TcpState;

#define TCP_DEFAULT_MSS (1200)
#define TCP_DEFAULT_RETRY_LIMIT (5)
#define TCP_DEFAULT_RETRY_TO_HS (100)

void tcps_init(TcpState *tcps, uint16_t localPort) {
    tcps->sequenceNumber = (uint32_t) rand();
    tcps->localPort = localPort;
    tcps->localMSS = TCP_DEFAULT_MSS;
    tcps->remotePort = 0;
    tcps->remoteAddr = 0;
    tcps->ackNumber = 0;
    tcps->window = ETHLIB_DEF_TCP_WINDOW_SIZE;
    tcps->txWin = tcpw_create(ETHLIB_DEF_TCP_WINDOW_SIZE);
    //tcps->rxWin = tcpw_create(ETHLIB_DEF_TCP_WINDOW_SIZE);
    tcps->retryLimit = TCP_DEFAULT_RETRY_LIMIT;
    tcps->retryTO = TCP_DEFAULT_RETRY_TO_HS;
    tcps->retries = 0;
    tcps->debug = false;
    tcps->userCb = NULL;
}

#define TCP_FETCH_STATE_FROM_CONNBLOCK(connBlock) ((TcpState *) (connBlock)->sieveLayer->tag.p)

void tcp_init_connection(ConnBlock *connBlock);

static void retry_to_connect_cb(Timer *timer, AlarmUserData user) {
    ConnBlock *connBlock = (ConnBlock *) user.ptr;
    tcp_init_connection(connBlock);
}

int tcp_receive_segment_cb(const Pckt *pckt, PcktSieveLayerTag tag) {
    TcpState *tcps = (TcpState *) tag.p; // pointer to state is stored in sieve layer tag
    TcpProps *tcpProps = HEADER_FETCH_PROPS(TcpProps, pckt->header); // fetch header

    TcpConnectionState beginState = tcps->connState;

    IPv4Props *ipProps = HEADER_FETCH_PROPS(IPv4Props, pckt->header->prev);
    uint16_t dataSize = ipProps->TotalLength - tcpProps->DataOffset - ETH_IPv4_HEADER_SIZE; // extract data size

    int ret = 0;
    bool willRetry = false; // latest operation is going to be reattempted

    // process incoming packet
    switch (tcps->connState) {
        case TCP_STATE_SYN_SENT: {
            if (!(tcpProps->Flags & TCP_FLAG_ACK)) {
                // TODO: if not acknowledged
            }
            if (tcpProps->Flags & TCP_FLAG_SYN) {
                tcps->ackNumber = tcpProps->SequenceNumber + 1;
            }

            // get MSS
            TcpOption *mss = tcp_get_option_by_kind(tcpProps->options, TCP_OPT_KIND_MSS);
            if (mss == NULL) {
                tcps->connState = TCP_STATE_CLOSED;
                if (tcps->retries < tcps->retryLimit) {
                    AlarmUserData alarmData;
                    alarmData.ptr = &tcps->connBlock;
                    timer_sched_rel(E.tmr, tcps->retryTO * 10000, retry_to_connect_cb, alarmData);
                    willRetry = true; // indicate that operation is going to be reattempted
                }
                break;
            }
            FETCH_WORD_H2N(&tcps->remoteMSS, mss->value);

            // send ACK
            int flags = TCP_FLAG_ACK;
            tcps->sequenceNumber++; // advance sequence number // TODO...
            tcp_send_segment(&tcps->connBlock, flags, NULL, NULL, 0);

            // step into next state
            tcps->connState = TCP_STATE_ESTAB;
            break;
        }
        case TCP_STATE_ESTAB:
            // if the other end tries to close down the connection
            if (tcpProps->Flags & TCP_FLAG_FIN) {
                // send FIN, ACK
                int flags = TCP_FLAG_FIN | TCP_FLAG_ACK;
                tcps->ackNumber++; // advance acknowledgement number // TODO...
                tcp_send_segment(&tcps->connBlock, flags, NULL, NULL, 0);

                // step into next state
                tcps->connState = TCP_STATE_LAST_ACK;
            } else if ((dataSize > 0) && (tcps->ackNumber == tcpProps->SequenceNumber)) { // incoming data segment with integrity in ack/seq
                // send acknowledge
                tcps->ackNumber += dataSize;
                tcp_send_segment(&tcps->connBlock, TCP_FLAG_ACK, NULL, NULL, 0);

                // process data
                if (tcps->userCb != NULL) {
                    PcktSieveLayerTag userTag = {0}; // TODO...
                    tcps->userCb(pckt, userTag);
                }
            } else if ((dataSize == 0) && (tcps->sequenceNumber == tcpProps->AcknowledgementNumber)) { // outgoing segment was acknowledged by peer
                // release segment from retransmit window
                tcpw_acknowledge(tcps->txWin, tcpProps->AcknowledgementNumber);
            }
            break;
        case TCP_STATE_LAST_ACK:
            if (tcpProps->Flags & TCP_FLAG_ACK) {
                tcps->connState = TCP_STATE_CLOSED;
                ret = SIEVE_LAYER_REMOVE_THIS; // if peer closed the connection, remove sieve layer
            }
            break;
        case TCP_STATE_CLOSED:
        default:
            break;
    }

    // print state
    if (tcps->debug && (beginState != tcps->connState)) {
        MSG("%s -> %s\n", TCP_STATE_NAMES[beginState], TCP_STATE_NAMES[tcps->connState]);
    }

    // maintain retry counter
    tcps->retries = willRetry ? (tcps->retries + 1) : 0;

    // release descriptor if not meaningful anymore
    if (ret == SIEVE_LAYER_REMOVE_THIS) {
        cbdt_release(E.cbdt, tcps->d); // release descriptor
    }

    return 0; // TODO
}

void tcp_init_connection(ConnBlock *connBlock) {
    /* send SYN with maximum segment size option */

    // TODO: check that our state is CLOSED

    TcpState *tcps = TCP_FETCH_STATE_FROM_CONNBLOCK(connBlock);

    // prepare MSS option
    TcpOption *mss = tcpopt_new(TCP_OPT_KIND_MSS, 2);
    FILL_WORD_H2N(mss->value, tcps->localMSS);

    // prepare flags
    int flags = TCP_FLAG_SYN;

    // send segment
    tcp_send_segment(connBlock, flags, mss, NULL, 0);

    // release option
    dynmem_free(mss);

    // step into SYN SENT state
    tcps->connState = TCP_STATE_SYN_SENT;
}

void tcp_bind(cbd d, ip4_addr remoteAddr, uint16_t remotePort) {
    ConnBlock connBlock;
    if (!cbdt_get_connection_block(E.cbdt, d, &connBlock)) {
        ERROR("Invalid CBD descriptor: '%d'!\n", d);
        return;
    }

    // store remote endpoint data
    TcpState *tcps = TCP_FETCH_STATE_FROM_CONNBLOCK(&connBlock);
    tcps->remoteAddr = remoteAddr;
    tcps->remotePort = remotePort;

    // initialize connection
    tcp_init_connection(&connBlock);
}

void tcp_debug(cbd d, bool debug) {
    ConnBlock connBlock;
    if (!cbdt_get_connection_block(E.cbdt, d, &connBlock)) {
        ERROR("Invalid CBD descriptor: '%d'!\n", d);
        return;
    }

    TcpState *tcps = TCP_FETCH_STATE_FROM_CONNBLOCK(&connBlock);
    tcps->debug = debug;
}

// ---------------------

cbd tcp_new_connblock(EthInterface *intf, ip4_addr ipAddr, uint16_t port, SieveCallBackFn cbFn) {
    ConnBlock tcpConnB;
    connb_init_defaults(&tcpConnB);

    ConnBlock ipConnB = ipv4_new_connblock(intf, ipAddr, NULL); // create new IPv4 connection block

    PcktSieveFilterCondition filtCond;
    packfiltcond_zero(&filtCond);
    TCP_PORT_TO_FILTCOND(&filtCond, port);

    PcktSieveLayerTag tag; // store TCP state into sieve layer's tag
    TcpState *tcpState = dynmem_alloc(sizeof(TcpState));
    tcps_init(tcpState, port);
    tcpState->userCb = cbFn;
    tag.p = tcpState;

    tcpConnB.sieveLayer = packsieve_new_layer(ipConnB.sieveLayer, &filtCond, false, filtTcp, tcp_receive_segment_cb, tag, ETH_TCP_PACKET_CLASS);
    ASSERT_NULL(tcpConnB.sieveLayer);
    tcpConnB.sieveLayer->connBReportFn = tcp_print_report;

    tcpConnB.sieve = &intf->sieve;
    tcpState->connBlock = tcpConnB; // also store connection block parameters

    // store connection block to CBDT
    cbd d = cbdt_alloc_new(E.cbdt, &tcpConnB);
    if (d == CBDT_ERR) { // on error free everything we have allocated before
        packsieve_remove_layer(tcpConnB.sieveLayer);
    }
    tcpState->d = d; // save descriptor
    return d;
}

int tcp_send_segment(const struct ConnBlock_ *connBlock, TcpFlag flags, TcpOption *opts, const uint8_t *data, uint32_t size) {
    // allocate headers
    PcktHeaderElement *tcpHeader = ALLOC_HEADER_ELEMENT(TcpProps);
    PcktHeaderElement *ipHeader = ALLOC_HEADER_ELEMENT(IPv4Props);
    PcktHeaderElement *ethHeader = ALLOC_HEADER_ELEMENT(EthernetProps);
    tcpHeader->next = NULL;
    tcpHeader->prev = ipHeader;
    ipHeader->next = tcpHeader;
    ipHeader->prev = ethHeader;
    ethHeader->next = ipHeader;
    ethHeader->prev = NULL;

    // prepare headers
    TcpProps *tcpProps = HEADER_FETCH_PROPS(TcpProps, tcpHeader);
    IPv4Props *ipProps = HEADER_FETCH_PROPS(IPv4Props, ipHeader);
    EthernetProps *ethProps = HEADER_FETCH_PROPS(EthernetProps, ethHeader);

    // get TCP state
    PcktSieveLayer *layer = connBlock->sieveLayer; // TCP layer
    TcpState *tcpState = (TcpState *) layer->tag.p;

    // fetch sieve layers and fill transmit headers
    tcpProps->SourcePort = tcpState->localPort;
    tcpProps->DestinationPort = tcpState->remotePort;
    tcpProps->SequenceNumber = tcpState->sequenceNumber;
    tcpProps->AcknowledgementNumber = tcpState->ackNumber;
    tcpProps->Flags = flags;
    tcpProps->Window = tcpState->window;
    tcpProps->Checksum = 0;
    tcpProps->UrgentPtr = 0;
    tcpProps->options = opts;

    // get options size
    uint32_t optSize = tcp_get_options_size(tcpProps->options);

    // common fields for packet assembly
    tcpProps->hdrInsFn = insert_tcp_header;
    tcpProps->headerSize = ETH_TCP_HEADER_SIZE + optSize;

    // IP
    ipv4_fill_props(ipProps, ETH_TCP_HEADER_SIZE + optSize + size, ETH_TCP_PACKET_CLASS,
                    connBlock->sieve->intf->ip, tcpState->remoteAddr);

    // Ethernet
    layer = layer->parent;
    if (tcpState->remoteAddr != 0xFFFFFFFF) {
        ArpCache *arpc = connBlock->sieve->intf->arpc;
        const ArpEntry *entry = arpc_get_ask(arpc, tcpState->remoteAddr);
        if (entry == NULL) {
            INFO("HW address cannot be ARP-ed, cannot send TCP segment!\n");
            goto release_resources; // YEAH, goto HERE!
        }
        memcpy(ethProps->destAddr, entry, ETH_HW_ADDR_LEN);
    }
    memcpy(ethProps->sourceAddr, connBlock->sieve->intf->mac, ETH_HW_ADDR_LEN);
    ethProps->length_type = ETH_IPv4_PACKET_CLASS;

    // common fields for packet assembly
    ethProps->hdrInsFn = insert_ethernet_header;
    ethProps->headerSize = ETH_ETHERNET_HEADER_SIZE;

    // -----------------------------

    Pckt cooked;
    cooked.payload = data;
    cooked.payloadSize = size;
    cooked.header = ethHeader;

    // NOT FILLED FIELDS
    cooked.headerSize = 0;
    cooked.time_s = 0;
    cooked.time_ns = 0;

    RawPckt raw;
    pckt_assemble(&raw, &cooked, connBlock->sieve->intf);

    ethinf_transmit(connBlock->sieve->intf, &raw);

    // advance TCP sequence number
    tcpState->sequenceNumber += size;

    // free headers
    release_resources:
    dynmem_free(tcpHeader);
    dynmem_free(ipHeader);
    dynmem_free(ethHeader);

    return 0;
}

uint32_t tcp_send(cbd d, const uint8_t *data, uint32_t size) {
    ConnBlock connBlock;
    if (!cbdt_get_connection_block(E.cbdt, d, &connBlock)) {
        ERROR("Invalid CBD descriptor: '%d'!\n", d);
        return 0;
    }

    TcpState *tcps = TCP_FETCH_STATE_FROM_CONNBLOCK(&connBlock);
    if (tcps->connState == TCP_STATE_ESTAB) { // can send only data if connection is ESTABLISHED
        uint32_t maxWinSize = tcpw_get_max_window_size(tcps->txWin);
        uint32_t txSize = MIN(maxWinSize, size); // limit size to the allocable size of the retransmit buffer
        TcpWindowSegment *winSeg = tcpw_store_segment(tcps->txWin, data, txSize, tcps->sequenceNumber); // store segment for possible retransmission
        tcp_send_segment(&connBlock, TCP_FLAG_ACK, NULL, winSeg->data, winSeg->size);
        return txSize;
    } else {
        return 0;
    }
}

void tcp_print_report(const ConnBlock *connBlock) {
    const PcktSieveLayer *sl = connBlock->sieveLayer;
    INFO("TCP port: %d", TCP_PORT_FROM_FILTCOND(&sl->filtCond));
}