#include "net.h"
#include <vector>
#include "pico/unique_id.h"
#include "pico/time.h"
#include "eth.h"
#include "arp.h"
#include "ipv4.h"
#include "w6300.h"
#include "debug_log.h"

static constexpr ipv4::ip4_addr IP_BROADCAST_SUBNET = {169, 254, 255, 255};
static constexpr uint16_t PICOMAP_PORT = __builtin_bswap16(28781);

static net_state state;
static w6300::socket_id raw_socket{0};
static net_handler msg_handler;
static std::vector<net_frame_callback> frame_callbacks;

void net_send_raw(std::span<const uint8_t> data) {
    dlog_if_slow("net_send_raw", 1000, [&]{
        w6300::send(raw_socket, data);
    });
}

static void handle_udp(std::span<const uint8_t> frame, span_writer& tx) {
    if (frame.size() < sizeof(ipv4::udp_header)) return;
    auto& pkt = *reinterpret_cast<const ipv4::udp_header*>(frame.data());

    if (pkt.dst_port != PICOMAP_PORT) return;
    if (!msg_handler) return;

    size_t udp_len = __builtin_bswap16(pkt.length);
    if (udp_len < 8) return;
    if (sizeof(eth::header) + pkt.ip.ip_total_len() < sizeof(ipv4::udp_header) + udp_len - 8) return;

    size_t payload_len = udp_len - 8;

    eth::mac_addr dst_mac = pkt.ip.eth.src;
    ipv4::ip4_addr dst_ip = pkt.ip.src;
    uint16_t dst_port = pkt.src_port;

    msg_handler(frame.subspan(sizeof(ipv4::udp_header), payload_len),
        [dst_mac, dst_ip, dst_port](std::span<const uint8_t> resp_data) {
            size_t ip_total = 20 + 8 + resp_data.size();
            size_t reply_len = sizeof(eth::header) + ip_total;
            uint8_t reply_buf[1514];
            if (reply_len > sizeof(reply_buf)) return;

            auto& rip = *reinterpret_cast<ipv4::header*>(reply_buf);
            rip.eth.dst = dst_mac;
            rip.eth.src = state.mac;
            rip.eth.ethertype = eth::ETH_IPV4;
            rip.ver_ihl = 0x45;
            rip.dscp_ecn = 0;
            rip.total_len = __builtin_bswap16(ip_total);
            rip.identification = 0;
            rip.flags_frag = 0;
            rip.ttl = 64;
            rip.protocol = 17;
            rip.checksum = 0;
            rip.src = state.ip;
            rip.dst = dst_ip;
            rip.checksum = ipv4::checksum(rip.ip_start(), 20);

            auto& rudp = *reinterpret_cast<ipv4::udp_header*>(reply_buf);
            rudp.src_port = PICOMAP_PORT;
            rudp.dst_port = dst_port;
            rudp.length = __builtin_bswap16(8 + resp_data.size());
            rudp.checksum = 0;

            memcpy(reply_buf + sizeof(ipv4::udp_header), resp_data.data(), resp_data.size());
            net_send_raw({reply_buf, reply_len});
        });
}

static bool mac_match(const eth::mac_addr& dst) {
    return dst == state.mac || dst == eth::MAC_BROADCAST;
}

static void process_frame(std::span<const uint8_t> frame, span_writer& tx) {
    if (frame.size() < sizeof(eth::header)) return;
    auto& eth_hdr = *reinterpret_cast<const eth::header*>(frame.data());

    if (!mac_match(eth_hdr.dst)) return;

    auto cbs = std::move(frame_callbacks);
    frame_callbacks.clear();
    for (auto& cb : cbs)
        cb(frame);

    switch (eth_hdr.ethertype) {
    case eth::ETH_ARP:
        arp::handle(frame, tx, state.mac, state.ip, net_send_raw);
        break;
    case eth::ETH_IPV4:
        ipv4::handle(frame, tx, state.mac, state.ip, IP_BROADCAST_SUBNET, net_send_raw, handle_udp);
        break;
    }
}

bool net_init() {
    w6300::init_spi();
    w6300::reset();
    w6300::init();
    if (!w6300::check()) return false;

    pico_unique_board_id_t uid;
    pico_get_unique_board_id(&uid);
    state.mac[0] = (uid.id[0] & 0xFC) | 0x02;
    state.mac[1] = uid.id[1];
    state.mac[2] = uid.id[2];
    state.mac[3] = uid.id[3];
    state.mac[4] = uid.id[4];
    state.mac[5] = uid.id[5];

    state.ip[0] = 169;
    state.ip[1] = 254;
    state.ip[2] = state.mac[4];
    state.ip[3] = state.mac[5];

    w6300::open_socket(raw_socket, w6300::protocol::macraw, w6300::sock_flag::none);
    w6300::set_interrupt_mask(w6300::ik_sock_0);

    return true;
}

const net_state& net_get_state() {
    return state;
}

void net_set_handler(net_handler handler) {
    msg_handler = std::move(handler);
}

void net_add_frame_callback(net_frame_callback cb) {
    frame_callbacks.push_back(std::move(cb));
}

void net_poll(std::span<uint8_t> tx) {
    if (!w6300::irq_pending) return;
    w6300::irq_pending = false;
    w6300::clear_interrupt(w6300::ik_int_all);
    static uint8_t rx_buf[1518];
    while (w6300::get_socket_recv_buf(raw_socket) > 0) {
        auto result = w6300::recv(raw_socket, std::span{rx_buf});
        if (!result) break;
        span_writer tx_writer(tx);
        process_frame({rx_buf, *result}, tx_writer);
    }
}