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

static void test_discovery_igmp(const responder& resp) {
    auto& ns = net_get_state();

    prepend_buffer<128> buf;
    igmp::prepend_query(buf, ns.mac, ns.ip, igmp::PICOMAP_DISCOVERY_GROUP);
    net_send_raw(buf.span());

    auto timer = std::make_shared<timer_handle>(nullptr);
    auto cb = std::make_shared<std::function<void(std::span<const uint8_t>)>>();
    *cb = [resp, timer, cb](std::span<const uint8_t> frame) {
        ipv4::ip4_addr group;
        if (!igmp::parse_report(frame, group)) {
            net_add_frame_callback(*cb);
            return;
        }
        if (group != igmp::PICOMAP_DISCOVERY_GROUP) {
            net_add_frame_callback(*cb);
            return;
        }
        dispatch_cancel_timer(*timer);
        resp.respond(ResponseTest{true, {"got IGMP report for " + ipv4::to_string(group)}});
    };
    net_add_frame_callback(*cb);

    *timer = dispatch_schedule_ms(5000, [resp]() {
        resp.respond(ResponseTest{false, {"no IGMP report within 5s"}});
    });
}

static void test_discovery_info(const responder& resp) {
    auto& ns = net_get_state();

    eth::mac_addr mcast_mac = igmp::mac_for_ip(igmp::PICOMAP_DISCOVERY_GROUP);
    static constexpr uint16_t PICOMAP_PORT = __builtin_bswap16(28781);

    prepend_buffer<1514> 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) {
        resp.respond(ResponseTest{false, {"encode RequestInfo failed"}});
        return;
    }
    buf.append(*encoded);
    size_t payload_len = *encoded;

    udp::prepend(buf, mcast_mac, ns.mac, ns.ip, igmp::PICOMAP_DISCOVERY_GROUP,
                 PICOMAP_PORT, PICOMAP_PORT, payload_len, 1);
    net_send_raw(buf.span());

    ipv4::ip4_addr our_ip = ns.ip;

    auto timer = std::make_shared<timer_handle>(nullptr);
    auto cb = std::make_shared<std::function<void(std::span<const uint8_t>)>>();
    *cb = [resp, our_ip, timer, 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) {
            net_add_frame_callback(*cb);
            return;
        }
        auto* ip = pb.consume<ipv4::header>();
        if (!ip || ip->protocol != 17) {
            net_add_frame_callback(*cb);
            return;
        }
        size_t options_len = ip->header_len() - sizeof(ipv4::header);
        if (options_len > 0 && !pb.skip(options_len)) {
            net_add_frame_callback(*cb);
            return;
        }
        auto* uhdr = pb.consume<udp::header>();
        if (!uhdr || uhdr->src_port != __builtin_bswap16(28781)) {
            net_add_frame_callback(*cb);
            return;
        }
        if (ip->src == our_ip) {
            net_add_frame_callback(*cb);
            return;
        }
        dispatch_cancel_timer(*timer);
        resp.respond(ResponseTest{true, {"got info response from " + ipv4::to_string(ip->src)}});
    };
    net_add_frame_callback(*cb);

    *timer = dispatch_schedule_ms(5000, [resp]() {
        resp.respond(ResponseTest{false, {"no info response within 5s"}});
    });
}

static void test_ping(const responder& resp, ipv4::ip4_addr dst_ip) {
    auto& ns = net_get_state();
    uint16_t ping_id = 0x1234;

    prepend_buffer<128> buf;
    icmp::prepend_echo_request(buf, ns.mac, ns.ip,
                               eth::MAC_BROADCAST, dst_ip, ping_id, 1);
    net_send_raw(buf.span());

    ipv4::ip4_addr our_ip = ns.ip;

    auto timer = std::make_shared<timer_handle>(nullptr);
    auto cb = std::make_shared<std::function<void(std::span<const uint8_t>)>>();
    *cb = [resp, ping_id, our_ip, timer, cb](std::span<const uint8_t> frame) {
        ipv4::ip4_addr src_ip;
        if (!icmp::parse_echo_reply(frame, src_ip, ping_id)) {
            net_add_frame_callback(*cb);
            return;
        }
        dispatch_cancel_timer(*timer);
        if (src_ip == our_ip) {
            resp.respond(ResponseTest{false, {"got reply from self: " + ipv4::to_string(src_ip)}});
            return;
        }
        resp.respond(ResponseTest{true, {"reply from " + ipv4::to_string(src_ip)}});
    };
    net_add_frame_callback(*cb);

    *timer = dispatch_schedule_ms(5000, [resp]() {
        resp.respond(ResponseTest{false, {"no reply from non-self host within 5s"}});
    });
}

static void test_ping_subnet(const responder& resp) {
    test_ping(resp, {169, 254, 255, 255});
}

static void test_ping_global(const responder& resp) {
    test_ping(resp, {255, 255, 255, 255});
}

using sync_test_fn = ResponseTest (*)(const responder&);
using async_test_fn = void (*)(const responder&);

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}},
};

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) {
    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(resp);
    it->second.async(resp);
    return std::nullopt;
}