#include "net.h"
#include <cstring>
#include "pico/unique_id.h"
#include "pico/time.h"
#include "w6300.h"
#include "debug_log.h"

static net_state state;
static w6300::socket_id raw_socket{0};

static constexpr uint16_t ETHERTYPE_ARP = 0x0806;
static constexpr uint16_t ETHERTYPE_IPV4 = 0x0800;
static constexpr uint8_t IP_PROTO_ICMP = 1;
static constexpr uint8_t ICMP_ECHO_REQUEST = 8;
static constexpr uint8_t ICMP_ECHO_REPLY = 0;
static constexpr uint16_t ARP_OP_REQUEST = 1;
static constexpr uint16_t ARP_OP_REPLY = 2;

static uint16_t read_u16(const uint8_t* p) { return (p[0] << 8) | p[1]; }

static void write_u16(uint8_t* p, uint16_t v) {
    p[0] = v >> 8;
    p[1] = v & 0xFF;
}

static uint16_t ip_checksum(const uint8_t* data, size_t len) {
    uint32_t sum = 0;
    for (size_t i = 0; i < len - 1; i += 2)
        sum += read_u16(data + i);
    if (len & 1)
        sum += data[len - 1] << 8;
    while (sum >> 16)
        sum = (sum & 0xFFFF) + (sum >> 16);
    return ~sum;
}

static bool mac_match(const uint8_t* dst) {
    static constexpr uint8_t broadcast[6] = {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
    return memcmp(dst, state.mac.data(), 6) == 0 ||
           memcmp(dst, broadcast, 6) == 0;
}

static bool ip_match(const uint8_t* dst) {
    return memcmp(dst, state.ip.data(), 4) == 0;
}

static void send_raw(const uint8_t* data, size_t len) {
    dlog_if_slow("send_raw", 1000, [&]{
        w6300::ip_address dummy = {};
        w6300::sendto(raw_socket, std::span<const uint8_t>{data, len}, dummy, w6300::port_num{0});
    });
}

static void handle_arp(const uint8_t* frame, size_t len) {
    if (len < 42) return;
    const uint8_t* arp = frame + 14;

    if (read_u16(arp) != 1) return;
    if (read_u16(arp + 2) != ETHERTYPE_IPV4) return;
    if (arp[4] != 6 || arp[5] != 4) return;

    if (read_u16(arp + 6) != ARP_OP_REQUEST) return;
    if (!ip_match(arp + 24)) return;

    uint8_t reply[42];
    memcpy(reply, frame + 6, 6);
    memcpy(reply + 6, state.mac.data(), 6);
    write_u16(reply + 12, ETHERTYPE_ARP);

    uint8_t* rarp = reply + 14;
    write_u16(rarp, 1);
    write_u16(rarp + 2, ETHERTYPE_IPV4);
    rarp[4] = 6;
    rarp[5] = 4;
    write_u16(rarp + 6, ARP_OP_REPLY);
    memcpy(rarp + 8, state.mac.data(), 6);
    memcpy(rarp + 14, state.ip.data(), 4);
    memcpy(rarp + 18, arp + 8, 6);
    memcpy(rarp + 24, arp + 14, 4);

    send_raw(reply, 42);
}

static void handle_icmp(const uint8_t* frame, size_t len) {
    const uint8_t* ip = frame + 14;
    size_t ip_hdr_len = (ip[0] & 0x0F) * 4;
    size_t ip_total_len = read_u16(ip + 2);

    if (14 + ip_total_len > len) return;
    if (ip[9] != IP_PROTO_ICMP) return;
    if (!ip_match(ip + 16)) return;

    const uint8_t* icmp = ip + ip_hdr_len;
    size_t icmp_len = ip_total_len - ip_hdr_len;
    if (icmp_len < 8) return;

    if (icmp[0] != ICMP_ECHO_REQUEST) return;

    uint8_t reply[1514];
    size_t reply_len = 14 + ip_total_len;
    if (reply_len > sizeof(reply)) return;

    memcpy(reply, frame + 6, 6);
    memcpy(reply + 6, state.mac.data(), 6);
    write_u16(reply + 12, ETHERTYPE_IPV4);

    uint8_t* rip = reply + 14;
    memcpy(rip, ip, ip_hdr_len);
    memcpy(rip + 12, ip + 16, 4);
    memcpy(rip + 16, ip + 12, 4);
    rip[8] = 64;
    memset(rip + 10, 0, 2);
    uint16_t ip_cksum = ip_checksum(rip, ip_hdr_len);
    write_u16(rip + 10, ip_cksum);

    uint8_t* ricmp = rip + ip_hdr_len;
    memcpy(ricmp, icmp, icmp_len);
    ricmp[0] = ICMP_ECHO_REPLY;
    memset(ricmp + 2, 0, 2);
    uint16_t icmp_cksum = ip_checksum(ricmp, icmp_len);
    write_u16(ricmp + 2, icmp_cksum);

    send_raw(reply, reply_len);
}

static void handle_ipv4(const uint8_t* frame, size_t len) {
    if (len < 34) return;
    const uint8_t* ip = frame + 14;
    if ((ip[0] >> 4) != 4) return;

    handle_icmp(frame, len);
}

static void process_frame(const uint8_t* frame, size_t len) {
    if (len < 14) return;

    if (!mac_match(frame)) return;

    uint16_t ethertype = read_u16(frame + 12);

    switch (ethertype) {
    case ETHERTYPE_ARP:
        handle_arp(frame, len);
        break;
    case ETHERTYPE_IPV4:
        handle_ipv4(frame, len);
        break;
    }
}

bool net_init() {
    w6300::init_spi();
    w6300::init_critical_section();
    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::port_num{0}, w6300::sock_flag::none);
    w6300::set_socket_io_mode(raw_socket, w6300::sock_io_mode::nonblock);

    return true;
}

const net_state& net_get_state() {
    return state;
}

void net_poll() {
    if (w6300::get_socket_recv_buf(raw_socket) == 0) return;
    static uint8_t rx_buf[1518];
    w6300::ip_address dummy_addr = {};
    w6300::port_num dummy_port{0};
    auto result = w6300::recvfrom(raw_socket, std::span{rx_buf}, dummy_addr, dummy_port);
    if (!result) return;
    process_frame(rx_buf, *result);
}