"""Official firmware images: validation and layout.

Only images listed in firmware/firmware.toml with a matching SHA-256 are
offered for flashing: boxes normally get new firmware over OTA, so the tool
only ever writes the official factory image.

The factory image is a merged ESP32-S3 image (bootloader at 0x0, partition
table at 0x8000, app at ota_0). App-only flashing takes the app out of it
and writes it to ota_0; the layout comes from the image's own partition
table, not from constants.
"""

from __future__ import annotations

import hashlib
import struct
import tomllib
from dataclasses import dataclass
from pathlib import Path

FIRMWARE_DIR = Path(__file__).resolve().parents[2] / "firmware"
MANIFEST = FIRMWARE_DIR / "firmware.toml"

PT_OFFSET = 0x8000
PT_SIZE = 0xC00
ESP32S3_CHIP_ID = 9
IMAGE_MAGIC = 0xE9
APP_DESC_MAGIC = 0xABCD5432
FLASH_SIZES = {0: "1MB", 1: "2MB", 2: "4MB", 3: "8MB", 4: "16MB", 5: "32MB"}


class ImageError(ValueError):
    pass


@dataclass(frozen=True)
class Partition:
    label: str
    type: int
    subtype: int
    offset: int
    size: int


@dataclass(frozen=True)
class AppInfo:
    project: str
    version: str
    built: str
    idf: str
    elf_sha: str       # hex


@dataclass(frozen=True)
class FirmwareImage:
    path: Path
    name: str
    notes: str
    sha256: str
    flash_size: str
    partitions: tuple[Partition, ...]
    app: AppInfo

    @property
    def data(self) -> bytes:
        return self.path.read_bytes()

    def partition(self, label: str) -> Partition:
        for p in self.partitions:
            if p.label == label:
                return p
        raise ImageError(f"partition {label!r} not in image")

    @property
    def partition_table(self) -> bytes:
        return self.data[PT_OFFSET:PT_OFFSET + PT_SIZE]

    def app_bytes(self) -> bytes:
        """The application image stored in the ota_0 region of the merged image."""
        # No stripping of trailing 0xFF: the image ends in a SHA-256 digest
        # that may legitimately end in 0xFF. The merged file ends with the app.
        ota0 = self.partition("ota_0")
        return self.data[ota0.offset:ota0.offset + ota0.size]

    def blank_otadata(self) -> bytes:
        """Erased otadata: the bootloader then boots ota_0."""
        return b"\xff" * self.partition("otadata").size


def parse_partition_table(raw: bytes) -> tuple[Partition, ...]:
    parts = []
    for i in range(0, len(raw) - 31, 32):
        entry = raw[i:i + 32]
        magic = entry[:2]
        if magic == b"\xaa\x50":
            _, ptype, sub, off, size = struct.unpack_from("<HBBII", entry)
            label = entry[12:28].split(b"\0")[0].decode("ascii", "replace")
            parts.append(Partition(label, ptype, sub, off, size))
        elif magic in (b"\xeb\xeb", b"\xff\xff"):   # MD5 entry / end
            break
        else:
            raise ImageError(f"corrupt partition table entry at 0x{PT_OFFSET + i:x}")
    return tuple(parts)


def _app_info(data: bytes, offset: int) -> AppInfo:
    if len(data) < offset + 0x20 + 256 or data[offset] != IMAGE_MAGIC:
        raise ImageError(f"no application image at 0x{offset:x}")
    d = data[offset + 0x20:offset + 0x20 + 256]
    if struct.unpack_from("<I", d)[0] != APP_DESC_MAGIC:
        raise ImageError("application has no ESP-IDF app descriptor")
    s = lambda o, n: d[o:o + n].split(b"\0")[0].decode("ascii", "replace")
    return AppInfo(project=s(48, 32), version=s(16, 32), built=f"{s(96, 16)} {s(80, 16)}",
                   idf=s(112, 32), elf_sha=d[144:176].hex())


def inspect_image(path: Path, name: str = "", notes: str = "") -> FirmwareImage:
    data = path.read_bytes()
    if len(data) < PT_OFFSET + PT_SIZE or data[0] != IMAGE_MAGIC:
        raise ImageError("not a merged ESP32 image (no bootloader at offset 0)")
    chip_id = struct.unpack_from("<H", data, 12)[0]
    if chip_id != ESP32S3_CHIP_ID:
        raise ImageError(f"image is for chip id {chip_id}, not ESP32-S3")
    flash_size = FLASH_SIZES.get(data[3] >> 4, "?")
    parts = parse_partition_table(data[PT_OFFSET:PT_OFFSET + PT_SIZE])
    labels = {p.label for p in parts}
    for need in ("nvs", "otadata", "ota_0"):
        if need not in labels:
            raise ImageError(f"partition table has no {need!r} partition")
    ota0 = next(p for p in parts if p.label == "ota_0")
    app = _app_info(data, ota0.offset)
    if app.project != "kota-sensorboard":
        raise ImageError(f"image is {app.project!r}, not kota-sensorboard")
    return FirmwareImage(path, name or path.name, notes, hashlib.sha256(data).hexdigest(),
                         flash_size, parts, app)


def official_images(manifest: Path = MANIFEST) -> list[FirmwareImage]:
    """Images from the manifest whose file exists and matches its SHA-256."""
    if not manifest.exists():
        return []
    entries = tomllib.loads(manifest.read_text(encoding="utf-8")).get("image", [])
    out = []
    for e in entries:
        path = manifest.parent / e["file"]
        if not path.exists():
            continue
        img = inspect_image(path, e.get("name", ""), e.get("notes", ""))
        if img.sha256 != e["sha256"].lower():
            raise ImageError(f"{path.name}: SHA-256 does not match the manifest - file changed or corrupt")
        out.append(img)
    return out