#include "dispatch.h"
#include <unordered_map>
#include "pico/stdlib.h"
#include "tusb.h"
#include "wire.h"
#include "usb_cdc.h"
#include "timer_queue.h"
#include "net.h"
#include "debug_log.h"
#include "hardware/sync.h"

static timer_queue timers;

void dispatch_init() {
    tusb_init();
    net_init();
    dlog("dispatch_init complete");
}

void dispatch_schedule_ms(uint32_t ms, std::function<void()> fn) {
    timers.schedule_ms(ms, std::move(fn));
}

[[noreturn]] void dispatch_run(std::span<const handler_entry> handlers) {
    std::unordered_map<int8_t, handler_fn> handler_map;
    for (auto& entry : handlers) {
        handler_map[entry.type_id] = entry.handle;
    }

    static usb_cdc usb;
    static static_vector<uint8_t, 256> usb_rx_buf;
    static std::array<uint8_t, 1514> tx_buf;

    net_set_handler([&](std::span<const uint8_t> payload, span_writer &out) -> msgpack::result<size_t> {
        auto msg = try_decode(payload.data(), payload.size());
        if (!msg) return 0;
        auto it = handler_map.find(msg->type_id);
        if (it == handler_map.end()) return 0;
        return it->second(msg->message_id, msg->payload, out);
    });

    while (true) {
        uint32_t save = save_and_disable_interrupts();

        dlog_if_slow("tud_task", 1000, [&]{ tud_task(); });
        dlog_if_slow("drain", 1000, [&]{ usb.drain(); });
        dlog_if_slow("timers", 1000, [&]{ timers.run(); });
        dlog_if_slow("net_poll", 1000, [&]{ net_poll(std::span{tx_buf}); });

        while (tud_cdc_available()) {
            uint8_t byte;
            if (tud_cdc_read(&byte, 1) != 1) break;

            usb_rx_buf.push_back(byte);

            auto msg = try_decode(usb_rx_buf);
            if (!msg) {
                if (usb_rx_buf.full()) usb_rx_buf.clear();
                continue;
            }

            auto it = handler_map.find(msg->type_id);
            if (it == handler_map.end()) { usb_rx_buf.clear(); continue; }
            span_writer out(tx_buf);
            auto resp = it->second(msg->message_id, msg->payload, out);
            usb_rx_buf.clear();
            if (!resp || *resp == 0) continue;
            size_t resp_len = *resp;
            if (resp_len <= usb.tx.free()) {
                usb.send(std::span<const uint8_t>{tx_buf.data(), resp_len});
                continue;
            }
            span_writer err_out(tx_buf);
            auto err = encode_response_into(err_out, msg->message_id,
                DeviceError{2, "response too large: " + std::to_string(resp_len)});
            if (err) usb.send(std::span<const uint8_t>{tx_buf.data(), *err});
        }

        __wfi();
        restore_interrupts(save);
    }
}