#include "handlers.h"
#include "pico/unique_id.h"
#include "pico/bootrom.h"
#include "hardware/flash.h"
#include "hardware/watchdog.h"
#include "boot/picobin.h"
#include "dispatch.h"
#include "net.h"
#include "debug_log.h"

static constexpr uint32_t XIP_BASE_ADDR = 0x10000000;
static constexpr uint32_t FLASH_SIZE = 2 * 1024 * 1024;
static constexpr uint32_t SLOT_A_OFFSET = 0x00000;
static constexpr uint32_t SLOT_B_OFFSET = 0x80000;
static boot_reason detected_boot_reason;

static void poke_watchdog() {
    watchdog_update();
    dispatch_schedule_ms(500, poke_watchdog);
}

void handlers_init() {
    auto val = static_cast<boot_reason>(watchdog_hw->scratch[0]);
    if (val == boot_reason::request_reboot || val == boot_reason::watchdog)
        detected_boot_reason = val;
    else
        detected_boot_reason = boot_reason::cold_boot;
    watchdog_hw->scratch[0] = static_cast<uint32_t>(boot_reason::watchdog);
    watchdog_enable(1000, true);
}

void handlers_start() {
    poke_watchdog();
}

std::optional<ResponsePICOBOOT> handle_picoboot(const responder&, const RequestPICOBOOT&) {
    dispatch_schedule_ms(100, []{ reset_usb_boot(0, 1); });
    return ResponsePICOBOOT{};
}

std::optional<ResponseInfo> handle_info(const responder&, const RequestInfo&) {
    ResponseInfo resp;
    pico_unique_board_id_t uid;
    pico_get_unique_board_id(&uid);
    std::copy(uid.id, uid.id + 8, resp.board_id.begin());
    auto& ns = net_get_state();
    resp.mac = ns.mac;
    resp.ip = ns.ip;
    resp.firmware_name = firmware_name;
    resp.boot = detected_boot_reason;
    resp.build_epoch = BUILD_EPOCH;
    return resp;
}

std::optional<ResponseLog> handle_log(const responder&, const RequestLog&) {
    ResponseLog resp;
    for (auto& e : dlog_drain())
        resp.entries.push_back(LogEntry{e.timestamp_us, std::move(e.message)});
    return resp;
}

std::optional<ResponseFlashErase> handle_flash_erase(const responder&, const RequestFlashErase& req) {
    if (req.addr < XIP_BASE_ADDR || req.addr + req.len > XIP_BASE_ADDR + FLASH_SIZE) {
        dlogf("flash erase: out of range %08lx+%lu",
              static_cast<unsigned long>(req.addr), static_cast<unsigned long>(req.len));
        return std::nullopt;
    }
    uint32_t offset = req.addr - XIP_BASE_ADDR;
    if (offset % FLASH_SECTOR_SIZE != 0 || req.len % FLASH_SECTOR_SIZE != 0 || req.len == 0) {
        dlogf("flash erase: bad alignment %08lx+%lu",
              static_cast<unsigned long>(req.addr), static_cast<unsigned long>(req.len));
        return std::nullopt;
    }
    flash_range_erase(offset, req.len);
    return ResponseFlashErase{};
}

std::optional<ResponseFlashWrite> handle_flash_write(const responder&, const RequestFlashWrite& req) {
    if (req.addr < XIP_BASE_ADDR || req.addr + req.data.size() > XIP_BASE_ADDR + FLASH_SIZE) {
        dlogf("flash write: out of range %08lx+%zu",
              static_cast<unsigned long>(req.addr), req.data.size());
        return std::nullopt;
    }
    uint32_t offset = req.addr - XIP_BASE_ADDR;
    if (offset % FLASH_PAGE_SIZE != 0 || req.data.size() % FLASH_PAGE_SIZE != 0 || req.data.empty()) {
        dlogf("flash write: bad alignment %08lx+%zu",
              static_cast<unsigned long>(req.addr), req.data.size());
        return std::nullopt;
    }
    flash_range_program(offset, req.data.data(), req.data.size());
    return ResponseFlashWrite{};
}

static SlotInfo scan_slot(uint32_t flash_offset) {
    SlotInfo info{};
    constexpr uint32_t scan_limit = 4096;
    auto* scan = reinterpret_cast<const uint32_t*>(XIP_BASE_ADDR + flash_offset);
    auto* scan_end = scan + scan_limit / 4;
    while (scan < scan_end && *scan != PICOBIN_BLOCK_MARKER_START) scan++;
    if (scan >= scan_end) return info;
    info.valid = true;

    auto* p = reinterpret_cast<const uint8_t*>(scan + 1);
    auto* end = reinterpret_cast<const uint8_t*>(scan + 64);
    while (p + 2 <= end) {
        uint8_t type = p[0];
        uint8_t size_words = p[1];
        if (type == PICOBIN_BLOCK_ITEM_2BS_LAST) break;
        uint32_t item_bytes = static_cast<uint32_t>(size_words) * 4;
        if (p + item_bytes > end) break;
        if (type == PICOBIN_BLOCK_ITEM_1BS_VERSION && size_words >= 2) {
            auto* words = reinterpret_cast<const uint16_t*>(p + 4);
            uint16_t minor = words[0];
            uint16_t major = words[1];
            info.version = (static_cast<uint32_t>(major) << 16) | minor;
        }
        p += item_bytes;
    }
    return info;
}

std::optional<ResponseFlashStatus> handle_flash_status(const responder&, const RequestFlashStatus&) {
    ResponseFlashStatus resp;
    boot_info_t bi;
    if (rom_get_boot_info(&bi))
        resp.boot_partition = bi.partition;
    else
        resp.boot_partition = -1;
    resp.slot_a = scan_slot(SLOT_A_OFFSET);
    resp.slot_b = scan_slot(SLOT_B_OFFSET);
    return resp;
}

std::optional<ResponseReboot> handle_reboot(const responder&, const RequestReboot&) {
    dispatch_schedule_ms(100, []{
        watchdog_hw->scratch[0] = static_cast<uint32_t>(boot_reason::request_reboot);
        watchdog_reboot(0, 0, 0);
    });
    return ResponseReboot{};
}