#include "flash.h"
#include <cstring>
#include "pico/sha256.h"
#include "boot/picobin.h"

namespace flash {
namespace {

constexpr uint32_t PICOBIN_MARKER_END = 0xab123579;

struct __attribute__((packed)) last_item {
    uint8_t type;
    uint16_t block_item_words;
    uint8_t pad;
    int32_t next_block_offset;
    uint32_t marker_end;
};

struct __attribute__((packed)) hash_def_header {
    uint8_t type;
    uint8_t size_words;
    uint8_t reserved;
    uint8_t hash_type;
    uint16_t block_words_to_hash;
    uint16_t pad;
};

struct __attribute__((packed)) load_map_header {
    uint8_t type;
    uint8_t size_words;
    uint8_t reserved;
    uint8_t flags_and_count;

    uint8_t count() const { return flags_and_count & 0x7f; }
    bool absolute() const { return flags_and_count & 0x80; }
};

struct load_map_entry {
    uint32_t storage_addr;
    uint32_t runtime_addr;
    uint32_t size;
};

struct parsed_item {
    uint8_t type;
    uint8_t size_words;
    const uint32_t* words;
};

struct parsed_block {
    const uint32_t* base;
    parsed_item items[16];
    uint8_t item_count;
    int32_t next_block_offset;
};

bool parse_block(const uint32_t* start, const uint32_t* limit, parsed_block& out) {
    if (start >= limit || *start != PICOBIN_BLOCK_MARKER_START) return false;
    out.base = start;
    out.item_count = 0;
    out.next_block_offset = 0;

    auto* w = start + 1;
    while (w + 3 < limit && out.item_count < 16) {
        uint8_t type = *w & 0xff;
        if (type == PICOBIN_BLOCK_ITEM_2BS_LAST) {
            auto* last = reinterpret_cast<const last_item*>(w);
            out.next_block_offset = last->next_block_offset;
            return last->marker_end == PICOBIN_MARKER_END;
        }
        uint8_t size_words = (*w >> 8) & 0xff;
        if (w + size_words > limit) return false;
        out.items[out.item_count++] = {type, size_words, w};
        w += size_words;
    }
    return false;
}

const parsed_item* find_item(const parsed_block& blk, uint8_t type) {
    for (uint8_t i = 0; i < blk.item_count; i++)
        if (blk.items[i].type == type) return &blk.items[i];
    return nullptr;
}

bool verify_hash(const parsed_block& last) {
    auto* lm_item = find_item(last, PICOBIN_BLOCK_ITEM_LOAD_MAP);
    auto* hd_item = find_item(last, PICOBIN_BLOCK_ITEM_1BS_HASH_DEF);
    auto* hv_item = find_item(last, PICOBIN_BLOCK_ITEM_HASH_VALUE);
    if (!lm_item || !hd_item || !hv_item) return false;
    if (hd_item->size_words < 2 || hv_item->size_words < 9) return false;

    auto* hd = reinterpret_cast<const hash_def_header*>(hd_item->words);
    if (hd->hash_type != PICOBIN_HASH_SHA256) return false;

    auto* lm = reinterpret_cast<const load_map_header*>(lm_item->words);
    auto* entries = reinterpret_cast<const load_map_entry*>(lm_item->words + 1);
    uint32_t lm_xip_addr = reinterpret_cast<uint32_t>(lm_item->words);

    pico_sha256_state_t sha;
    if (pico_sha256_try_start(&sha, SHA256_BIG_ENDIAN, false) != PICO_OK) return false;

    for (uint8_t i = 0; i < lm->count(); i++) {
        uint32_t storage_addr = entries[i].storage_addr;
        uint32_t size = entries[i].size;
        if (lm->absolute()) size -= entries[i].runtime_addr;
        if (storage_addr == 0) {
            pico_sha256_update_blocking(&sha, reinterpret_cast<const uint8_t*>(&size), 4);
        } else {
            if (!lm->absolute()) storage_addr += lm_xip_addr;
            pico_sha256_update_blocking(&sha, reinterpret_cast<const uint8_t*>(storage_addr), size);
        }
    }

    pico_sha256_update_blocking(&sha,
        reinterpret_cast<const uint8_t*>(last.base),
        static_cast<size_t>(hd->block_words_to_hash) * 4);

    sha256_result_t result;
    pico_sha256_finish(&sha, &result);

    auto* expected = reinterpret_cast<const uint8_t*>(hv_item->words + 1);
    return memcmp(result.bytes, expected, 32) == 0;
}

}

slot scan(uint32_t flash_offset) {
    slot info{};
    constexpr uint32_t scan_limit = 4096;
    constexpr uint32_t slot_size = 512 * 1024;
    auto* slot_base = reinterpret_cast<const uint32_t*>(FLASH_BASE + flash_offset);
    auto* slot_end = reinterpret_cast<const uint32_t*>(FLASH_BASE + flash_offset + slot_size);

    auto* s = slot_base;
    auto* s_end = slot_base + scan_limit / 4;
    while (s < s_end && *s != PICOBIN_BLOCK_MARKER_START) s++;
    if (s >= s_end) return info;

    parsed_block start;
    if (!parse_block(s, slot_end, start)) return info;

    for (uint8_t i = 0; i < start.item_count; i++) {
        if (start.items[i].type == PICOBIN_BLOCK_ITEM_1BS_VERSION && start.items[i].size_words >= 2)
            info.version = start.items[i].words[1];
    }

    auto* cur = &start;
    parsed_block next;
    while (cur->next_block_offset != 0) {
        auto* np = reinterpret_cast<const uint32_t*>(
            reinterpret_cast<const uint8_t*>(cur->base) + cur->next_block_offset);
        if (np <= slot_base || np >= slot_end) return info;
        if (np == start.base) break;
        if (!parse_block(np, slot_end, next)) return info;
        cur = &next;
    }

    info.valid = true;
    info.hash_ok = verify_hash(*cur);
    return info;
}

}