#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"

// TODO: - retransmission hiányzik
// TODO: - állapotváltás esetleg a küldés előtt? (külcsönös kizárás és társai...)
// TODO: - ne próbáljunk két kapcsolatot nyitni ugyanarra a portra...
// TODO: - ez az "ok" változó nem túl jó...
// TODO: - retry connection nincs kész

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

static inline cbd tcp_new_connblock_filtcond(EthInterface *intf, ip4_addr ipAddr, uint16_t port, SieveCallBackFn cbFn,
                               const PcktSieveFilterCondition *filtCond);

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

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;

    TCP_EXTRACT_FILTCOND(filtCond);

    bool references_us = (local_port == tcpProps->DestinationPort) && // it's our local port
                         (((remote_port == tcpProps->SourcePort) && // it's the current connections remote port
                           (remote_addr == ipProps->SourceIPAddr)) || // the packet has come from the known origin
                          (remote_addr == 0)); // the latter conditions only matter if the control block is connected to a remote station

    return ipProps->Protocol == ETH_TCP_PACKET_CLASS && references_us;
}

static inline TcpState * tcps_create() {
    TcpState * tcps = dynmem_alloc(sizeof(TcpState));
    tcps->txWin = tcpw_create(ETHLIB_DEF_TCP_WINDOW_SIZE);

    return tcps;
}

static inline void tcps_remove_and_cleanup(const TcpState * tcps) {
    if (tcps == NULL) {
        return;
    }

    dynmem_free(tcps->txWin);
    dynmem_free(tcps);
}

#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->connState = TCP_STATE_CLOSED;

    tcps->localPort = localPort;
    tcps->localMSS = TCP_DEFAULT_MSS;
    tcps->remotePort = 0;
    tcps->remoteAddr = 0;

    tcps->txSeqNum = (uint32_t) rand();
    tcps->rxAckNum = 0;

    tcps->localWindow = ETHLIB_DEF_TCP_WINDOW_SIZE;
    tcpw_set_seqnum_offset(tcps->txWin, tcps->txSeqNum);

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

static inline cbd tcps_create_based_on_listening(const TcpState *source, uint16_t remotePort, ip4_addr remoteAddr, ConnBlock *connBlock) {
    // construct filter condition
    PcktSieveFilterCondition filtCond;
    TCP_LOCAL_PORT_TO_FILTCOND(&filtCond, source->localPort);
    TCP_REMOTE_PORT_TO_FILTCOND(&filtCond, remotePort);
    TCP_REMOTE_ADDR_TO_FILTCOND(&filtCond, remoteAddr);

    // create connblock
    EthInterface *intf = source->connBlock.sieve->intf; // retrieve interface

    cbd d = tcp_new_connblock_filtcond(intf, IPv4_IF_ADDR, source->localPort, source->userCb, &filtCond);
    if (!cbdt_get_connection_block(E.cbdt, d, connBlock)) {
        ERROR("Invalid CBD descriptor: '%d'!\n", d);
        return CBDT_ERR;
    }

    // alter fields
    TcpState *target = TCP_FETCH_STATE_FROM_CONNBLOCK(connBlock);
    target->txSeqNum = source->txSeqNum;
    target->rxAckNum = source->rxAckNum;
    target->remoteAddr = remoteAddr;
    target->remotePort = remotePort;
    target->connState = TCP_STATE_ESTAB;
    tcpw_set_seqnum_offset(target->txWin, target->txSeqNum);

    return d;
}

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: { /* SYN segment sent to the server */

            bool ok = false; // operations were fine

            // if no ACK received, then it's likely an error
            if (!(tcpProps->Flags & TCP_FLAG_ACK)) {
                break;
            }

            // SYN and ACK received
            if ((tcpProps->Flags & TCP_FLAG_SYN) && (tcpProps->Flags & TCP_FLAG_ACK)) {
                // store rxAckNum
                tcps->rxAckNum = tcpProps->SequenceNumber + 1;

                // get MSS
                TcpOption *mss = tcp_get_option_by_kind(tcpProps->options, TCP_OPT_KIND_MSS);
                if (mss != NULL) { // if MSS is included in the segment
                    FETCH_WORD_H2N(&tcps->remoteMSS, mss->value); // store remote Maximum Segment Size
                    ok = true;
                } else { // if MSS is NOT included in the received segment...
                    tcps->connState = TCP_STATE_CLOSED; // close the initiating connection
                    if (tcps->retries < tcps->retryLimit) { // schedule a connection retry
                        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
                    }
                }

                // if previous processing concluded without an error
                if (ok) {
                    // send ACK
                    int flags = TCP_FLAG_ACK;
                    tcps->txSeqNum++; // advance sequence number // TODO...
                    tcp_send_segment(&tcps->connBlock, flags, NULL, NULL, 0);

                    // save filter values
                    PcktSieveFilterCondition *filtCond = &tcps->connBlock.sieveLayer->filtCond;
                    TCP_REMOTE_PORT_TO_FILTCOND(filtCond, tcpProps->SourcePort);
                    TCP_REMOTE_ADDR_TO_FILTCOND(filtCond, ipProps->SourceIPAddr);

                    // step into next state
                    tcps->connState = TCP_STATE_ESTAB;
                }

            }

            break;
        }

        case TCP_STATE_LISTEN: { /* Listening for incoming connections */
            if (tcpProps->Flags & TCP_FLAG_SYN) { // no ACK, just SYN

                bool ok = false;

                // store rxAckNum
                tcps->rxAckNum = tcpProps->SequenceNumber + 1;

                // get MSS
                TcpOption *mss = tcp_get_option_by_kind(tcpProps->options, TCP_OPT_KIND_MSS);
                if (mss != NULL) { // if MSS is included in the segment
                    FETCH_WORD_H2N(&tcps->remoteMSS, mss->value); // store remote Maximum Segment Size
                    ok = true;
                }

                // if previous processing concluded without an error
                if (ok) {
                    // send ACK
                    int flags = TCP_FLAG_ACK | TCP_FLAG_SYN;
                    tcps->txSeqNum++; // advance sequence number // TODO...

                    // create connection based on the listening connection
                    ConnBlock connB;
                    tcps_create_based_on_listening(tcps, tcpProps->SourcePort, ipProps->SourceIPAddr, &connB);
                    tcp_send_segment(&connB, flags, NULL, NULL, 0);

                    TcpState *target = TCP_FETCH_STATE_FROM_CONNBLOCK(&connB);
                    target->txSeqNum++;
                }

            }

            break;
        }

        case TCP_STATE_ESTAB: /* Connection established */
            // 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->rxAckNum++; // 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->rxAckNum == tcpProps->SequenceNumber)) { // incoming data segment with integrity in ack/seq
                // send acknowledge
                tcps->rxAckNum += 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->txSeqNum == 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: /* last ACK received */
            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: /* connection is CLOSED */
        default:
            break;
    }

    // save remote window size
    tcps->remoteWindow = tcpProps->Window;

    // 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
        tcps_remove_and_cleanup(tcps); // release TCP state
    }

    return ret;
}

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 station data
    TcpState *tcps = TCP_FETCH_STATE_FROM_CONNBLOCK(&connBlock);
    tcps->remoteAddr = remoteAddr;
    tcps->remotePort = remotePort;

    // 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_listen(cbd d) {
    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);
    if ((tcps->userCb != NULL) && (tcps->connState == TCP_STATE_CLOSED)) {
        tcps->connState = TCP_STATE_LISTEN;
    }
}

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

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

static inline cbd tcp_new_connblock_filtcond(EthInterface *intf, ip4_addr ipAddr, uint16_t port, SieveCallBackFn cbFn, const PcktSieveFilterCondition *filtCond) {
    ConnBlock tcpConnB;
    connb_init_defaults(&tcpConnB);

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

    PcktSieveLayerTag tag; // store TCP state into sieve layer's tag
    TcpState *tcpState = tcps_create();
    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;
}

cbd tcp_new_connblock(EthInterface *intf, ip4_addr ipAddr, uint16_t port, SieveCallBackFn cbFn) {
    PcktSieveFilterCondition filtCond;
    packfiltcond_zero(&filtCond);
    TCP_LOCAL_PORT_TO_FILTCOND(&filtCond, port);
    TCP_REMOTE_PORT_TO_FILTCOND(&filtCond, 0);
    TCP_REMOTE_ADDR_TO_FILTCOND(&filtCond, 0);

    return tcp_new_connblock_filtcond(intf, ipAddr, port, cbFn, &filtCond);
}

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->txSeqNum;
    tcpProps->AcknowledgementNumber = tcpState->rxAckNum;
    tcpProps->Flags = flags;
    tcpProps->Window = tcpState->localWindow;
    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->txSeqNum += 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(tcps->remoteWindow, MIN(maxWinSize, size)); // limit size to the allocatable size of the retransmit buffer
        tcpw_store_segment(tcps->txWin, data, txSize, tcps->txSeqNum); // store segment for possible retransmission
        tcp_send_segment(&connBlock, TCP_FLAG_ACK, NULL, data, txSize);
        return txSize;
    } else {
        return 0;
    }
}

void tcp_print_report(const ConnBlock *connBlock) {
    const PcktSieveLayer *sl = connBlock->sieveLayer;
    TcpState *tcps = TCP_FETCH_STATE_FROM_CONNBLOCK(connBlock);
    if (tcps->connState != TCP_STATE_ESTAB) {
        INFO("TCP port: %d - %s", tcps->localPort, TCP_STATE_NAMES[tcps->connState]);
    } else {
        INFO("TCP :%d -- %u.%u.%u.%u:%d - %s", tcps->localPort, EXPLODE_IPv4(tcps->remoteAddr), tcps->remotePort, TCP_STATE_NAMES[tcps->connState]);
    }
}