"""
Pure-Python ESP32 NVS partition binary generator.

Generates a binary-compatible NVS partition for a Vesper device identity.
No ESP-IDF toolchain required on the server.

NVS partition layout (ESP32 NVS format v2):
- Partition size: 0x5000 (20480 bytes) = 5 pages
- Page size: 4096 bytes
- Page structure:
    Offset 0x000 - 0x01F : Page header (32 bytes)
    Offset 0x020 - 0x03F : Entry state bitmap (32 bytes, 2 bits per slot)
    Offset 0x040 - 0xFFF : Entry storage (120 slots × 32 bytes each)

Entry state bitmap: 2 bits per entry
  11 = empty
  10 = written (active)
  00 = erased

Page header (32 bytes):
  uint32 page_state    (0xFFFFFFFE = active)
  uint32 sequence_number
  uint8  version       (0xFE = v2)
  uint8  reserved[19]
  uint32 crc32         (of bytes 4..27)

Entry (32 bytes):
  uint8  ns_index    (namespace index, 0 = namespace entry itself)
  uint8  type        (0x01=uint8, 0x02=uint16, 0x04=uint32, 0x08=uint64, 0x21=string, 0x41=blob)
  uint8  span        (number of 32-byte slots this entry occupies)
  uint8  chunk_index (0xFF for non-blob)
  uint32 crc32       (of the entry header bytes 0..3 and data, excluding the crc field itself)
  char   key[16]     (null-terminated, max 15 chars + null)
  <data> [8 bytes for primitives, or inline for short strings]

For strings:
  - If len <= 8 bytes (incl. null): fits in the data field of the same entry (span=1)
  - Longer strings: data follows in subsequent 32-byte "data entries" (span = 1 + ceil(strlen+1, 32))
  - The entry header data field contains: uint16 data_size, uint16 reserved=0xFFFF, uint32 crc32_of_data
"""

import struct
import binascii
from typing import List, Tuple


NVS_PAGE_SIZE = 4096
NVS_PARTITION_SIZE = 0x5000  # 20480 bytes = 5 pages
NVS_ENTRY_SIZE = 32
NVS_ENTRY_COUNT = 126  # entries per page (first 3 slots are header + bitmap)

NVS_PAGE_STATE_ACTIVE = 0xFFFFFFFE
NVS_PAGE_VERSION = 0xFE

ENTRY_STATE_WRITTEN = 0b10  # 2 bits
ENTRY_STATE_EMPTY   = 0b11  # 2 bits (erased flash)

ENTRY_TYPE_NAMESPACE = 0x01  # used for namespace entries (uint8)
ENTRY_TYPE_STRING    = 0x21


def _crc32(data: bytes) -> int:
    # ESP-IDF uses 0xFFFFFFFF as the initial CRC seed (matches esp_rom_crc32_le)
    return binascii.crc32(data, 0xFFFFFFFF) & 0xFFFFFFFF


def _page_header_crc(seq: int, version: int) -> int:
    """CRC covers bytes 4..27 of the page header (seq + version + reserved)."""
    buf = struct.pack("<IB", seq, version) + b"\xFF" * 19
    return _crc32(buf)


def _entry_crc(ns_index: int, entry_type: int, span: int, chunk_index: int,
               key: bytes, data: bytes) -> int:
    """CRC covers the entry minus the 4-byte crc field at offset 4..7."""
    header_no_crc = struct.pack("BBBB", ns_index, entry_type, span, chunk_index)
    return _crc32(header_no_crc + key + data)


def _pack_entry(ns_index: int, entry_type: int, span: int, chunk_index: int,
                key: str, data: bytes) -> bytes:
    key_bytes = key.encode("ascii").ljust(16, b"\x00")[:16]
    data_bytes = data.ljust(8, b"\xFF")[:8]
    crc = _entry_crc(ns_index, entry_type, span, chunk_index, key_bytes, data_bytes)
    return struct.pack("BBBBI", ns_index, entry_type, span, chunk_index, crc) + key_bytes + data_bytes


def _bitmap_set_written(bitmap: bytearray, slot_index: int) -> None:
    """Mark a slot as written (10) in the entry state bitmap."""
    bit_pos = slot_index * 2
    byte_idx = bit_pos // 8
    bit_off  = bit_pos % 8
    # Clear both bits for this slot (set to 00 then OR in 10)
    bitmap[byte_idx] &= ~(0b11 << bit_off)
    bitmap[byte_idx] |=  (ENTRY_STATE_WRITTEN << bit_off)


def _build_namespace_entry(ns_name: str, ns_index: int) -> Tuple[bytes, int]:
    """Build the namespace declaration entry. ns_index is the assigned namespace id (1-based)."""
    data = struct.pack("<B", ns_index) + b"\xFF" * 7
    entry = _pack_entry(
        ns_index=0,
        entry_type=ENTRY_TYPE_NAMESPACE,
        span=1,
        chunk_index=0xFF,
        key=ns_name,
        data=data,
    )
    return entry, 1  # consumes 1 slot


def _build_string_entry(ns_index: int, key: str, value: str) -> Tuple[bytes, int]:
    """Build a string entry. May span multiple 32-byte slots for long strings."""
    value_bytes = value.encode("utf-8") + b"\x00"  # null-terminated
    value_len = len(value_bytes)

    # Pad to multiple of 32
    padded_len = ((value_len + 31) // 32) * 32
    value_padded = value_bytes.ljust(padded_len, b"\xFF")

    span = 1 + (padded_len // 32)

    # Data field in the header entry: uint16 data_size, uint16 0xFFFF, uint32 crc_of_data
    data_crc = _crc32(value_bytes)
    header_data = struct.pack("<HHI", value_len, 0xFFFF, data_crc)

    entry = _pack_entry(
        ns_index=ns_index,
        entry_type=ENTRY_TYPE_STRING,
        span=span,
        chunk_index=0xFF,
        key=key,
        data=header_data,
    )

    # Append data chunks (each 32 bytes)
    full_entry = entry + value_padded
    return full_entry, span


def _build_page(entries: List[bytes], slot_counts: List[int], seq: int = 0) -> bytes:
    """Assemble a full 4096-byte NVS page."""
    # Build entry storage area
    storage = bytearray(NVS_ENTRY_COUNT * NVS_ENTRY_SIZE)  # all 0xFF (erased)
    storage[:] = b"\xFF" * len(storage)
    bitmap = bytearray(b"\xFF" * 32)  # all slots empty (11 bits)

    slot = 0
    for entry_bytes, span in zip(entries, slot_counts):
        entry_offset = slot * NVS_ENTRY_SIZE
        storage[entry_offset:entry_offset + len(entry_bytes)] = entry_bytes
        for s in range(span):
            _bitmap_set_written(bitmap, slot + s)
        slot += span

    # Page header
    header_crc = _page_header_crc(seq, NVS_PAGE_VERSION)
    header = struct.pack(
        "<IIBI19sI",
        NVS_PAGE_STATE_ACTIVE,
        seq,
        NVS_PAGE_VERSION,
        0,  # padding
        b"\xFF" * 19,
        header_crc,
    )
    # Trim header to exactly 32 bytes
    header = struct.pack("<I", NVS_PAGE_STATE_ACTIVE)
    header += struct.pack("<I", seq)
    header += struct.pack("<B", NVS_PAGE_VERSION)
    header += b"\xFF" * 19
    header += struct.pack("<I", header_crc)
    assert len(header) == 32, f"Header size mismatch: {len(header)}"

    page = header + bytes(bitmap) + bytes(storage)
    assert len(page) == NVS_PAGE_SIZE, f"Page size mismatch: {len(page)}"
    return page


def generate(serial_number: str, hw_family: str, hw_revision: str, legacy: bool = False) -> bytes:
    """Generate a 0x5000-byte NVS partition binary for a Vesper device.

    serial_number: full SN string e.g. 'BSVSPR-26C13X-STD01R-X7KQA'
    hw_family: board family e.g. 'vesper-standard', 'vesper-plus'
    hw_revision: hardware revision string e.g. '1.0'
    legacy: if True, writes old key names expected by legacy firmware (pre-new-schema):
              device_uid, hw_type, hw_version
            if False (default), writes new key names:
              serial, hw_family, hw_revision

    Returns raw bytes ready to flash at 0x9000.
    """
    ns_index = 1  # first (and only) namespace

    # Build entries for namespace "device_id"
    ns_entry, ns_span = _build_namespace_entry("device_id", ns_index)
    if legacy:
        uid_entry, uid_span = _build_string_entry(ns_index, "device_uid",  serial_number)
        hwt_entry, hwt_span = _build_string_entry(ns_index, "hw_type",     hw_family.lower())
        hwv_entry, hwv_span = _build_string_entry(ns_index, "hw_version",  hw_revision)
    else:
        uid_entry, uid_span = _build_string_entry(ns_index, "serial",      serial_number)
        hwt_entry, hwt_span = _build_string_entry(ns_index, "hw_family",   hw_family.lower())
        hwv_entry, hwv_span = _build_string_entry(ns_index, "hw_revision", hw_revision)

    entries = [ns_entry, uid_entry, hwt_entry, hwv_entry]
    spans   = [ns_span,  uid_span,  hwt_span,  hwv_span]

    page0 = _build_page(entries, spans, seq=0)

    # Remaining pages are blank (erased flash = 0xFF)
    blank_page = b"\xFF" * NVS_PAGE_SIZE
    remaining_pages = (NVS_PARTITION_SIZE // NVS_PAGE_SIZE) - 1

    return page0 + blank_page * remaining_pages