//
// Created by epagris on 2022.11.07..
//

#include "udp_connblock.h"

#include <stddef.h>

#include "../../utils.h"
#include "../../dynmem.h"
#include "../../pckt_assembler.h"
#include "../../global_state.h"

static bool filtUdp(const PcktSieveFilterCondition *filtCond, const PcktProps *contProps, const PcktProps *ownProps, EthInterface *intf) {
    IPv4Props *ipProps = (IPv4Props *) contProps;
    UdpProps *udpProps = (UdpProps *) ownProps;

    return ipProps->Protocol == ETH_UDP_PACKET_CLASS && (UDP_PORT_FROM_FILTCOND(filtCond) == udpProps->DestinationPort);
}

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

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

    PcktSieveFilterCondition filtCond;
    packfiltcond_zero(&filtCond);
    UDP_PORT_TO_FILTCOND(&filtCond, port);
    PcktSieveLayerTag tag;
    tag.u = 0;
    udpConnB.sieveLayer = packsieve_new_layer(ipConnB.sieveLayer, &filtCond, false, filtUdp, cbFn, tag, ETH_UDP_PACKET_CLASS);
    ASSERT_NULL(udpConnB.sieveLayer);
    udpConnB.sieveLayer->connBReportFn = udp_print_report;

    udpConnB.sieve = &intf->sieve;

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

int udp_sendto(cbd d, const uint8_t *data, uint32_t size, ip4_addr addr, uint16_t port) {
    return udp_sendto_arg(d, data, size, addr, port, 0);
}

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

    // allocate headers
    PcktHeaderElement *udpHeader = ALLOC_HEADER_ELEMENT(UdpProps);
    PcktHeaderElement *ipHeader = ALLOC_HEADER_ELEMENT(IPv4Props);
    PcktHeaderElement *ethHeader = ALLOC_HEADER_ELEMENT(EthernetProps);
    udpHeader->next = NULL;
    udpHeader->prev = ipHeader;
    ipHeader->next = udpHeader;
    ipHeader->prev = ethHeader;
    ethHeader->next = ipHeader;
    ethHeader->prev = NULL;

    // prepare headers
    UdpProps *udpProps = HEADER_FETCH_PROPS(UdpProps, udpHeader);
    IPv4Props *ipProps = HEADER_FETCH_PROPS(IPv4Props, ipHeader);
    EthernetProps *ethProps = HEADER_FETCH_PROPS(EthernetProps, ethHeader);

    // fetch sieve layers and fill transmit headers
    PcktSieveLayer *layer = connBlock.sieveLayer; // UDP layer
    udpProps->SourcePort = UDP_PORT_FROM_FILTCOND(&layer->filtCond);
    udpProps->DestinationPort = port;
    udpProps->Length = ETH_UDP_HEADER_SIZE + size;
    udpProps->Checksum = 0;

    // common fields for packet assembly
    udpProps->hdrInsFn = insert_udp_header;
    udpProps->headerSize = ETH_UDP_HEADER_SIZE;

    // IP
    layer = layer->parent;
    ipv4_fill_props(ipProps, ETH_UDP_HEADER_SIZE + size, ETH_UDP_PACKET_CLASS,
                    connBlock.sieve->intf->ip, addr);

    // Ethernet
    layer = layer->parent;
    bool isMulticast = ipv4_is_multicast_address(addr);
    if ((!isMulticast) && (addr != 0xFFFFFFFF)) { // unicast
        ArpCache *arpc = connBlock.sieve->intf->arpc;
        const ArpEntry *entry = arpc_get_ask(arpc, addr);
        if (entry == NULL) {
            INFO("HW address cannot be ARP-ed, cannot send UDP datagram!\n");
            goto release_resources; // YEAH, goto HERE!
        }
        memcpy(ethProps->destAddr, entry, ETH_HW_ADDR_LEN);
    } else if (isMulticast) { // multicast
        ethmc_from_ipmc(ethProps->destAddr, addr);
    } else { // unicast
        memset(ethProps->destAddr, 0xFF, ETH_HW_ADDR_LEN); // broadcast
    }
    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);
    raw.ext.tx.arg = arg; // assign argument
    raw.ext.tx.txTsCb = connBlock.sieveLayer->txTsCb; // assign with TX timestamp callback

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

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

    // release payload
    // dynmem_free(raw.payload);

    return size; // TODO...
}

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