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

using mac_addr = std::array<uint8_t, 6>;
using ip4_addr = std::array<uint8_t, 4>;

struct __attribute__((packed)) eth_header {
    mac_addr dst;
    mac_addr src;
    uint16_t ethertype;
};
static_assert(sizeof(eth_header) == 14);

struct __attribute__((packed)) arp_packet {
    eth_header eth;
    uint16_t htype;
    uint16_t ptype;
    uint8_t hlen;
    uint8_t plen;
    uint16_t oper;
    mac_addr sha;
    ip4_addr spa;
    mac_addr tha;
    ip4_addr tpa;
};
static_assert(sizeof(arp_packet) == 42);

struct __attribute__((packed)) ipv4_header {
    eth_header eth;
    uint8_t ver_ihl;
    uint8_t dscp_ecn;
    uint16_t total_len;
    uint16_t identification;
    uint16_t flags_frag;
    uint8_t ttl;
    uint8_t protocol;
    uint16_t checksum;
    ip4_addr src;
    ip4_addr dst;

    size_t ip_header_len() const { return (ver_ihl & 0x0F) * 4; }
    size_t ip_total_len() const { return __builtin_bswap16(total_len); }
    const uint8_t* ip_start() const { return reinterpret_cast<const uint8_t*>(&ver_ihl); }
    uint8_t* ip_start() { return reinterpret_cast<uint8_t*>(&ver_ihl); }
};
static_assert(sizeof(ipv4_header) == 34);

struct __attribute__((packed)) icmp_echo {
    uint8_t type;
    uint8_t code;
    uint16_t checksum;
    uint16_t id;
    uint16_t seq;
};
static_assert(sizeof(icmp_echo) == 8);

static constexpr uint16_t ETH_ARP = __builtin_bswap16(0x0806);
static constexpr uint16_t ETH_IPV4 = __builtin_bswap16(0x0800);
static constexpr uint16_t ARP_HTYPE_ETH = __builtin_bswap16(1);
static constexpr uint16_t ARP_PTYPE_IPV4 = __builtin_bswap16(0x0800);
static constexpr uint16_t ARP_OP_REQUEST = __builtin_bswap16(1);
static constexpr uint16_t ARP_OP_REPLY = __builtin_bswap16(2);
static constexpr mac_addr MAC_BROADCAST = {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
static constexpr ip4_addr IP_BROADCAST_ALL = {255, 255, 255, 255};
static constexpr ip4_addr IP_BROADCAST_SUBNET = {169, 254, 255, 255};

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

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

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

static bool ip_match(const ip4_addr& dst) {
    return dst == state.ip;
}

static bool ip_match_or_broadcast(const ip4_addr& dst) {
    return ip_match(dst) || dst == IP_BROADCAST_ALL || dst == IP_BROADCAST_SUBNET;
}

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

static void handle_arp(const uint8_t* frame, size_t len) {
    if (len < sizeof(arp_packet)) return;
    auto& pkt = *reinterpret_cast<const arp_packet*>(frame);

    if (pkt.htype != ARP_HTYPE_ETH) return;
    if (pkt.ptype != ARP_PTYPE_IPV4) return;
    if (pkt.hlen != 6 || pkt.plen != 4) return;
    if (pkt.oper != ARP_OP_REQUEST) return;
    if (!ip_match(pkt.tpa)) return;

    arp_packet reply = {};
    reply.eth.dst = pkt.eth.src;
    reply.eth.src = state.mac;
    reply.eth.ethertype = ETH_ARP;
    reply.htype = ARP_HTYPE_ETH;
    reply.ptype = ARP_PTYPE_IPV4;
    reply.hlen = 6;
    reply.plen = 4;
    reply.oper = ARP_OP_REPLY;
    reply.sha = state.mac;
    reply.spa = state.ip;
    reply.tha = pkt.sha;
    reply.tpa = pkt.spa;

    send_raw(&reply, sizeof(reply));
}

static void handle_icmp(const uint8_t* frame, size_t len) {
    auto& ip = *reinterpret_cast<const ipv4_header*>(frame);
    size_t ip_hdr_len = ip.ip_header_len();
    size_t ip_total = ip.ip_total_len();

    if (sizeof(eth_header) + ip_total > len) return;
    if (ip.protocol != 1) return;
    if (!ip_match_or_broadcast(ip.dst)) return;

    auto& icmp = *reinterpret_cast<const icmp_echo*>(frame + sizeof(eth_header) + ip_hdr_len);
    size_t icmp_len = ip_total - ip_hdr_len;
    if (icmp_len < sizeof(icmp_echo)) return;
    if (icmp.type != 8) return;

    uint8_t reply_buf[1514];
    size_t reply_len = sizeof(eth_header) + ip_total;
    if (reply_len > sizeof(reply_buf)) return;

    memcpy(reply_buf, frame, reply_len);
    auto& rip = *reinterpret_cast<ipv4_header*>(reply_buf);
    rip.eth.dst = ip.eth.src;
    rip.eth.src = state.mac;
    rip.src = state.ip;
    rip.dst = ip.src;
    rip.ttl = 64;
    rip.checksum = 0;
    rip.checksum = ip_checksum(rip.ip_start(), ip_hdr_len);

    auto& ricmp = *reinterpret_cast<icmp_echo*>(reply_buf + sizeof(eth_header) + ip_hdr_len);
    ricmp.type = 0;
    ricmp.checksum = 0;
    ricmp.checksum = ip_checksum(&ricmp, icmp_len);

    send_raw(reply_buf, reply_len);
}

static void handle_ipv4(const uint8_t* frame, size_t len) {
    if (len < sizeof(ipv4_header)) return;
    auto& ip = *reinterpret_cast<const ipv4_header*>(frame);
    if ((ip.ver_ihl >> 4) != 4) return;

    switch (ip.protocol) {
    case 1:
        handle_icmp(frame, len);
        break;
    }
}

static void process_frame(const uint8_t* frame, size_t len) {
    if (len < sizeof(eth_header)) return;
    auto& eth = *reinterpret_cast<const eth_header*>(frame);

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

    switch (eth.ethertype) {
    case ETH_ARP:
        handle_arp(frame, len);
        break;
    case ETH_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);
}