#include <stdio.h>
#include <string.h>
#include <vector>
#include "pico/stdlib.h"
#include "pico/bootrom.h"
#include "msgpackpp.h"
#include "halfsiphash.h"
#include "static_vector.h"

static constexpr uint8_t hash_key[8] = {};

struct ResponseBOOTSEL {
    static constexpr int8_t ext_id = 1;
    auto as_tuple() const { return std::tie(); }
    auto as_tuple() { return std::tie(); }
};

struct RequestBOOTSEL {
    static constexpr int8_t ext_id = 2;
    auto as_tuple() const { return std::tie(); }
    auto as_tuple() { return std::tie(); }
};

struct Envelope {
    static constexpr int8_t ext_id = 0;
    uint32_t checksum;
    std::vector<uint8_t> payload;
    auto as_tuple() const { return std::tie(checksum, payload); }
    auto as_tuple() { return std::tie(checksum, payload); }
};

static std::vector<uint8_t> pack_envelope(const std::vector<uint8_t> &payload) {
    uint32_t checksum = halfsiphash::hash32(payload.data(), payload.size(), hash_key);
    msgpackpp::packer p;
    p.pack(Envelope{checksum, payload});
    return p.get_payload();
}

static void send_bytes(const std::vector<uint8_t> &data) {
    for (auto b : data) {
        putchar(b);
    }
    stdio_flush();
}

template <typename T>
static void send_message(const T &msg) {
    msgpackpp::packer inner;
    inner.pack(msg);
    auto envelope = pack_envelope(inner.get_payload());
    send_bytes(envelope);
}

static int8_t try_decode(const static_vector<uint8_t, 256> &buf) {
    msgpackpp::parser p(buf.data(), static_cast<int>(buf.size()));

    Envelope env;
    if (!msgpackpp::unpack(p, env)) return -1;

    uint32_t expected = halfsiphash::hash32(env.payload.data(), env.payload.size(), hash_key);
    if (env.checksum != expected) return -1;

    msgpackpp::parser inner(env.payload.data(), static_cast<int>(env.payload.size()));
    if (!inner.is_ext()) return -1;
    auto ext = inner.get_ext();
    if (!ext) return -1;

    return std::get<0>(*ext);
}

int main() {
    stdio_init_all();

    static static_vector<uint8_t, 256> rx_buf;

    while (true) {
        int c = getchar_timeout_us(100000);
        if (c == PICO_ERROR_TIMEOUT) continue;

        rx_buf.push_back(static_cast<uint8_t>(c));

        int8_t msg_type = try_decode(rx_buf);
        if (msg_type < 0) {
            if (rx_buf.full()) rx_buf.clear();
            continue;
        }

        rx_buf.clear();

        switch (msg_type) {
        case RequestBOOTSEL::ext_id:
            send_message(ResponseBOOTSEL{});
            sleep_ms(100);
            reset_usb_boot(0, 0);
            break;
        }
    }
}