#include "test_handlers.h"
#include <cstring>
#include <unordered_map>
#include "pico/stdlib.h"
#include "pico/time.h"
#include "handlers.h"
#include "net.h"
#include "icmp.h"
#include "igmp.h"
#include "udp.h"
#include "parse_buffer.h"
#include "prepend_buffer.h"

static constexpr uint16_t PING_ECHO_ID = 0x1234;
static constexpr uint16_t PING_RATE_ECHO_ID = 0x5678;

struct peer_info {
    eth::mac_addr mac;
    ipv4::ip4_addr ip;
};

struct discovery_data {
    void (*on_found)(const peer_info&) = nullptr;
    void (*on_timeout)() = nullptr;
};

struct ping_rate_data {
    peer_info peer;
    uint16_t target;
    uint16_t pipeline;
    uint16_t payload_len;
    uint16_t sent;
    uint16_t received;
    uint32_t start_us;
};

struct discovery_igmp_test {};
struct discovery_info_test {
    discovery_data discovery;
};
struct ping_subnet_test {};
struct ping_global_test {};
struct packet_rate_test {
    discovery_data discovery;
    ping_rate_data rate;
};
struct byte_rate_test {
    discovery_data discovery;
    ping_rate_data rate;
};

// One test runs at a time; in_flight gates that. active_* let shared
// primitive callbacks find the running test's sub-state.
struct test_state {
    bool in_flight = false;
    responder resp;
    timer_handle timer = nullptr;
    net::frame_cb_handle frame_cb = nullptr;

    discovery_data* active_discovery = nullptr;
    ping_rate_data* active_rate = nullptr;

    discovery_igmp_test discovery_igmp;
    discovery_info_test discovery_info;
    ping_subnet_test ping_subnet;
    ping_global_test ping_global;
    packet_rate_test packet_rate;
    byte_rate_test byte_rate;
};

static test_state ts;

static void test_end(const ResponseTest& result) {
    if (ts.timer) { dispatch_cancel_timer(ts.timer); ts.timer = nullptr; }
    if (ts.frame_cb) { net::remove_frame_callback(ts.frame_cb); ts.frame_cb = nullptr; }
    ts.active_discovery = nullptr;
    ts.active_rate = nullptr;
    ts.resp.respond(result);
    ts.in_flight = false;
}

// When a callback fires, its dispatcher (net or timer_queue) has already
// removed the node; the matching ts.timer / ts.frame_cb is stale. Callbacks
// that self-consume must null that handle before calling test_end.

static bool discover_reply_cb(std::span<const uint8_t> frame) {
    parse_buffer pb(frame);
    auto* eth_hdr = pb.consume<eth::header>();
    if (!eth_hdr || eth_hdr->ethertype != eth::ETH_IPV4) return false;
    auto* ip = pb.consume<ipv4::header>();
    if (!ip || ip->protocol != 17) return false;
    size_t options_len = ip->header_len() - sizeof(ipv4::header);
    if (options_len > 0 && !pb.skip(options_len)) return false;
    auto* uhdr = pb.consume<udp::header>();
    if (!uhdr || uhdr->src_port != PICOMAP_PORT_BE) return false;
    if (ip->src == net::get_state().ip) return false;
    dispatch_cancel_timer(ts.timer);
    ts.timer = nullptr;
    ts.frame_cb = nullptr;
    auto cont = ts.active_discovery ? ts.active_discovery->on_found : nullptr;
    ts.active_discovery = nullptr;
    peer_info peer{eth_hdr->src, ip->src};
    if (cont) cont(peer);
    return true;
}

static void discover_timeout_cb() {
    net::remove_frame_callback(ts.frame_cb);
    ts.frame_cb = nullptr;
    ts.timer = nullptr;
    auto cont = ts.active_discovery ? ts.active_discovery->on_timeout : nullptr;
    ts.active_discovery = nullptr;
    if (cont) cont();
}

static void discover_peer(discovery_data& d,
                          void (*found)(const peer_info&), void (*timeout)()) {
    d.on_found = found;
    d.on_timeout = timeout;
    ts.active_discovery = &d;

    const auto& ns = net::get_state();
    eth::mac_addr mcast_mac = igmp::mac_for_ip(PICOMAP_DISCOVERY_GROUP);

    prepend_buffer<4096> buf;
    uint8_t* payload = buf.payload_ptr();
    span_writer out(payload, 1024);
    RequestInfo req_msg;
    auto encoded = encode_response_into(out, 0xFFFF, req_msg);
    if (!encoded) {
        ts.active_discovery = nullptr;
        timeout();
        return;
    }
    buf.append(*encoded);

    udp::prepend(buf, mcast_mac, ns.mac, ns.ip, PICOMAP_DISCOVERY_GROUP,
                 PICOMAP_PORT_BE, PICOMAP_PORT_BE, *encoded, 1);

    ts.frame_cb = net::add_frame_callback(discover_reply_cb);
    ts.timer = dispatch_schedule_ms(5000, discover_timeout_cb);

    net::send_raw(buf.span());
}

static bool igmp_report_cb(std::span<const uint8_t> frame) {
    ipv4::ip4_addr group;
    if (!igmp::parse_report(frame, group)) return false;
    if (group != PICOMAP_DISCOVERY_GROUP) return false;
    ts.frame_cb = nullptr;
    test_end({true, {"got IGMP report for " + ipv4::to_string(group)}});
    return true;
}

static void igmp_timeout_cb() {
    ts.timer = nullptr;
    test_end({false, {"no IGMP report within 5s"}});
}

static void test_discovery_igmp() {
    const auto& ns = net::get_state();
    prepend_buffer<4096> buf;
    igmp::prepend_query(buf, ns.mac, ns.ip, PICOMAP_DISCOVERY_GROUP);

    ts.frame_cb = net::add_frame_callback(igmp_report_cb);
    ts.timer = dispatch_schedule_ms(5000, igmp_timeout_cb);

    net::send_raw(buf.span());
}

static void info_found(const peer_info& peer) {
    test_end({true, {"got info response from " + ipv4::to_string(peer.ip)}});
}

static void info_timeout() {
    test_end({false, {"no info response within 5s"}});
}

static void test_discovery_info() {
    discover_peer(ts.discovery_info.discovery, info_found, info_timeout);
}

static bool ping_reply_cb(std::span<const uint8_t> frame) {
    ipv4::ip4_addr src_ip;
    if (!icmp::parse_echo_reply(frame, src_ip, PING_ECHO_ID)) return false;
    ts.frame_cb = nullptr;
    if (src_ip == net::get_state().ip)
        test_end({false, {"got reply from self: " + ipv4::to_string(src_ip)}});
    else
        test_end({true, {"reply from " + ipv4::to_string(src_ip)}});
    return true;
}

static void ping_timeout_cb() {
    ts.timer = nullptr;
    test_end({false, {"no reply from non-self host within 5s"}});
}

static void start_ping(ipv4::ip4_addr dst_ip) {
    const auto& ns = net::get_state();
    prepend_buffer<4096> buf;
    icmp::prepend_echo_request(buf, ns.mac, ns.ip,
                               eth::MAC_BROADCAST, dst_ip, PING_ECHO_ID, 1);
    ts.frame_cb = net::add_frame_callback(ping_reply_cb);
    ts.timer = dispatch_schedule_ms(5000, ping_timeout_cb);
    net::send_raw(buf.span());
}

static void test_ping_subnet() { start_ping({169, 254, 255, 255}); }
static void test_ping_global() { start_ping({255, 255, 255, 255}); }

static size_t ping_rate_frame_size() {
    return sizeof(eth::header) + sizeof(ipv4::header) + sizeof(icmp::echo)
         + ts.active_rate->payload_len;
}

static void ping_rate_send_one() {
    const auto& ns = net::get_state();
    auto& r = *ts.active_rate;
    prepend_buffer<4096> buf;
    if (r.payload_len > 0)
        memset(buf.append(r.payload_len), 0xAA, r.payload_len);
    icmp::prepend_echo_request(buf, ns.mac, ns.ip,
                               r.peer.mac, r.peer.ip, PING_RATE_ECHO_ID,
                               r.sent + 1, r.payload_len);
    net::send_raw(buf.span());
    r.sent++;
}

static bool ping_rate_reply_cb(std::span<const uint8_t> frame) {
    ipv4::ip4_addr src_ip;
    if (!icmp::parse_echo_reply(frame, src_ip, PING_RATE_ECHO_ID)) return false;
    if (src_ip == net::get_state().ip) return false;

    auto& r = *ts.active_rate;
    r.received++;
    if (r.received >= r.target) {
        uint32_t elapsed_us = time_us_32() - r.start_us;
        uint32_t elapsed_ms = elapsed_us / 1000;
        uint32_t pps = static_cast<uint32_t>(
            static_cast<uint64_t>(r.received) * 1000000 / elapsed_us);
        uint64_t total_bytes = static_cast<uint64_t>(r.received) * 2 * ping_rate_frame_size();
        uint32_t kbps = static_cast<uint32_t>(total_bytes * 1000 / elapsed_us);
        char msg[128];
        snprintf(msg, sizeof(msg),
                 "%u rt in %lu ms, %lu pps, %lu bytes, %lu KB/s",
                 r.received, static_cast<unsigned long>(elapsed_ms),
                 static_cast<unsigned long>(pps),
                 static_cast<unsigned long>(total_bytes),
                 static_cast<unsigned long>(kbps));
        ts.frame_cb = nullptr;
        test_end({true, {msg}});
        return true;
    }

    if (r.sent < r.target)
        ping_rate_send_one();
    return false;
}

static void ping_rate_timeout_cb() {
    ts.timer = nullptr;
    auto& r = *ts.active_rate;
    uint32_t elapsed_us = time_us_32() - r.start_us;
    char msg[64];
    snprintf(msg, sizeof(msg), "timeout after %u/%u rt in %lu ms",
             r.received, r.sent,
             static_cast<unsigned long>(elapsed_us / 1000));
    test_end({false, {msg}});
}

static void ping_rate_found(const peer_info& peer) {
    auto& r = *ts.active_rate;
    r.peer = peer;
    r.sent = 0;
    r.received = 0;
    r.start_us = time_us_32();

    ts.frame_cb = net::add_frame_callback(ping_rate_reply_cb);
    ts.timer = dispatch_schedule_ms(10000, ping_rate_timeout_cb);

    for (uint16_t i = 0; i < r.pipeline && r.sent < r.target; i++)
        ping_rate_send_one();
}

static void ping_rate_no_peer() {
    test_end({false, {"no peer found"}});
}

static void start_ping_rate(discovery_data& d, ping_rate_data& r,
                            uint16_t target, uint16_t payload_len, uint16_t pipeline) {
    r.target = target;
    r.payload_len = payload_len;
    r.pipeline = pipeline;
    ts.active_rate = &r;
    discover_peer(d, ping_rate_found, ping_rate_no_peer);
}

static void test_packet_rate() {
    start_ping_rate(ts.packet_rate.discovery, ts.packet_rate.rate, 8192, 0, 8);
}
static void test_byte_rate() {
    start_ping_rate(ts.byte_rate.discovery, ts.byte_rate.rate, 2048, 1400, 8);
}

using sync_test_fn = ResponseTest (*)();
using async_test_fn = void (*)();

struct test_entry {
    sync_test_fn sync;
    async_test_fn async;
};

static const std::unordered_map<std::string_view, test_entry> tests = {
    {"discovery_igmp", {nullptr, test_discovery_igmp}},
    {"discovery_info", {nullptr, test_discovery_info}},
    {"ping_subnet",    {nullptr, test_ping_subnet}},
    {"ping_global",    {nullptr, test_ping_global}},
    {"packet_rate",    {nullptr, test_packet_rate}},
    {"byte_rate",      {nullptr, test_byte_rate}},
};

std::optional<ResponseListTests> handle_list_tests(const responder&, const RequestListTests&) {
    ResponseListTests resp;
    for (const auto& [name, _] : tests)
        resp.names.emplace_back(name);
    return resp;
}

std::optional<ResponseTest> handle_test(const responder& resp, const RequestTest& req) {
    if (ts.in_flight)
        return ResponseTest{false, {"test already running"}};
    auto it = tests.find(req.name);
    if (it == tests.end())
        return ResponseTest{false, {"unknown test: " + req.name}};
    if (it->second.sync)
        return it->second.sync();

    ts.in_flight = true;
    ts.resp = resp;
    it->second.async();
    return std::nullopt;
}