"""Firmware images: validation, integrity and layout.

Two origins:
  * factory   - bundled with the tool (firmware/firmware.toml), SHA-256 locked;
  * developer - a newer build the firmware team sends to technicians,
                imported into a per-user library (~/KotaTechnicianTool/firmware).
Two kinds:
  * merged - bootloader at 0x0 + partition table + app (like new_system.bin);
             App-only or Full flashing;
  * app    - the application binary only (build/kota-sensorboard.bin);
             App-only flashing, laid out like the factory image.

Every image is checked structurally (ESP32-S3, kota-sensorboard app
descriptor, segment checksum and the appended SHA-256) so a truncated or
corrupted file is refused before anything is written.
"""

from __future__ import annotations

import datetime as dt
import hashlib
import json
import shutil
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"
LIBRARY_DIR = Path.home() / "KotaTechnicianTool" / "firmware"
LIBRARY_INDEX = "library.json"

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"}
_MONTHS = {m: i for i, m in enumerate(
    ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"], 1)}


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          # "Jun 29 2026 11:03:32" as in the app descriptor
    idf: str
    elf_sha: str        # hex

    @property
    def built_at(self) -> dt.datetime | None:
        try:
            mon, day, year, hms = self.built.split()
            h, m, s = (int(x) for x in hms.split(":"))
            return dt.datetime(int(year), _MONTHS[mon], int(day), h, m, s)
        except (ValueError, KeyError):
            return None


@dataclass(frozen=True)
class FirmwareImage:
    path: Path
    name: str
    notes: str
    origin: str                      # "factory" | "developer"
    kind: str                        # "merged" | "app"
    sha256: str
    flash_size: str
    partitions: tuple[Partition, ...]
    partition_table: bytes           # expected table on the box
    app_offset: int                  # where the app starts inside the file
    app_length: int                  # exact app image length (from its structure)
    app: AppInfo

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

    @property
    def is_developer(self) -> bool:
        return self.origin == "developer"

    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")

    def app_bytes(self) -> bytes:
        return self.data[self.app_offset:self.app_offset + self.app_length]

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

    def label(self) -> str:
        tag = "factory" if self.origin == "factory" else "developer build"
        return f"{self.app.version} built {self.app.built} [{self.app.elf_sha[:9]}] - {tag}"


# ---------------------------------------------------------------- parsing

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_image_length(data: bytes, offset: int = 0) -> int:
    """Validate an ESP app image at `offset`; return its exact length.

    Checks the segment checksum and, when the header says so, the appended
    SHA-256 - the same checks the bootloader performs before booting it.
    """
    if len(data) < offset + 24 or data[offset] != IMAGE_MAGIC:
        raise ImageError(f"no firmware image at 0x{offset:x}")
    segments = data[offset + 1]
    chip_id = struct.unpack_from("<H", data, offset + 12)[0]
    if chip_id != ESP32S3_CHIP_ID:
        raise ImageError(f"image is for chip id {chip_id}, not ESP32-S3")
    hash_appended = data[offset + 23] == 1
    pos, checksum = offset + 24, 0xEF
    for _ in range(segments):
        if pos + 8 > len(data):
            raise ImageError("image is truncated (segment header missing)")
        _, seg_len = struct.unpack_from("<II", data, pos)
        pos += 8
        if pos + seg_len > len(data):
            raise ImageError("image is truncated (segment data missing)")
        for b in data[pos:pos + seg_len]:
            checksum ^= b
        pos += seg_len
    pos += 15 - ((pos - offset) % 16)          # pad so the checksum ends a 16-byte block
    if pos >= len(data) or data[pos] != checksum:
        raise ImageError("image checksum is wrong - the file is corrupted")
    pos += 1
    if hash_appended:
        if pos + 32 > len(data):
            raise ImageError("image is truncated (SHA-256 missing)")
        if hashlib.sha256(data[offset:pos]).digest() != data[pos:pos + 32]:
            raise ImageError("image SHA-256 does not match its content - the file is corrupted")
        pos += 32
    return pos - offset


def _app_info(data: bytes, offset: int) -> AppInfo:
    d = data[offset + 0x20:offset + 0x20 + 256]
    if len(d) < 256 or struct.unpack_from("<I", d)[0] != APP_DESC_MAGIC:
        raise ImageError("firmware has no ESP-IDF app descriptor")
    s = lambda o, n: d[o:o + n].split(b"\0")[0].decode("ascii", "replace")
    info = 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())
    if info.project != "kota-sensorboard":
        raise ImageError(f"firmware is {info.project!r}, not kota-sensorboard")
    return info


def inspect_image(path: Path, name: str = "", notes: str = "", origin: str = "factory",
                  reference: FirmwareImage | None = None) -> FirmwareImage:
    """Validate a merged or app-only image. App-only images need a `reference`
    (the factory image) that provides the flash layout."""
    data = path.read_bytes()
    if len(data) >= PT_OFFSET + PT_SIZE and data[0] == IMAGE_MAGIC and data[PT_OFFSET:PT_OFFSET + 2] == b"\xaa\x50":
        app_image_length(data, 0)              # bootloader must be intact
        table = data[PT_OFFSET:PT_OFFSET + PT_SIZE]
        parts = parse_partition_table(table)
        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")
        kind, app_off, flash_size = "merged", ota0.offset, FLASH_SIZES.get(data[3] >> 4, "?")
    elif data[:1] == bytes([IMAGE_MAGIC]):
        if reference is None:
            raise ImageError("app-only image needs the factory image as layout reference")
        kind, app_off = "app", 0
        parts, table, flash_size = reference.partitions, reference.partition_table, reference.flash_size
        ota0 = next(p for p in parts if p.label == "ota_0")
    else:
        raise ImageError("not an ESP32 firmware file")
    length = app_image_length(data, app_off)
    if length > ota0.size:
        raise ImageError(f"firmware ({length} bytes) does not fit the ota_0 partition ({ota0.size} bytes)")
    app = _app_info(data, app_off)
    return FirmwareImage(path, name or path.name, notes, origin, kind, hashlib.sha256(data).hexdigest(),
                         flash_size, parts, table, app_off, length, app)


# ------------------------------------------------------------- catalogue

def factory_images(manifest: Path = MANIFEST) -> list[FirmwareImage]:
    """Bundled images whose file exists and matches the SHA-256 in the manifest."""
    if not manifest.exists():
        return []
    out = []
    for e in tomllib.loads(manifest.read_text(encoding="utf-8")).get("image", []):
        path = manifest.parent / e["file"]
        if not path.exists():
            continue
        img = inspect_image(path, e.get("name", ""), e.get("notes", ""), "factory")
        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


class FirmwareLibrary:
    """Developer builds imported by the technician (persist across tool versions)."""

    def __init__(self, root: Path = LIBRARY_DIR):
        self.root = root

    def _index(self) -> list[dict]:
        f = self.root / LIBRARY_INDEX
        return json.loads(f.read_text(encoding="utf-8")) if f.exists() else []

    def _save(self, entries: list[dict]) -> None:
        self.root.mkdir(parents=True, exist_ok=True)
        (self.root / LIBRARY_INDEX).write_text(json.dumps(entries, indent=2), encoding="utf-8")

    def images(self, reference: FirmwareImage | None) -> tuple[list[FirmwareImage], list[str]]:
        """(valid images, problems) - a changed or broken file is reported, not offered."""
        imgs, problems = [], []
        for e in self._index():
            path = self.root / e["file"]
            try:
                img = inspect_image(path, e.get("name", ""), e.get("notes", ""), "developer", reference)
                if img.sha256 != e["sha256"]:
                    raise ImageError("file changed since it was imported")
                imgs.append(img)
            except (OSError, ImageError) as exc:
                problems.append(f"{e.get('original', e['file'])}: {exc}")
        return imgs, problems

    def import_file(self, src: Path, reference: FirmwareImage | None, notes: str = "") -> FirmwareImage:
        img = inspect_image(src, origin="developer", reference=reference)   # validate before copying
        entries = self._index()
        existing = next((e for e in entries if e["sha256"] == img.sha256), None)
        if existing:
            return inspect_image(self.root / existing["file"], existing["name"], existing["notes"],
                                 "developer", reference)
        self.root.mkdir(parents=True, exist_ok=True)
        fname = f"{img.app.elf_sha[:9]}_{img.sha256[:8]}.bin"
        shutil.copyfile(src, self.root / fname)
        name = f"Developer build {img.app.version} ({img.app.built})"
        entries.append({"file": fname, "sha256": img.sha256, "name": name, "original": src.name,
                        "imported": dt.datetime.now().isoformat(timespec="seconds"), "notes": notes})
        self._save(entries)
        return inspect_image(self.root / fname, name, notes, "developer", reference)

    def remove(self, sha256: str) -> None:
        entries = self._index()
        keep = [e for e in entries if e["sha256"] != sha256]
        for e in entries:
            if e["sha256"] == sha256:
                (self.root / e["file"]).unlink(missing_ok=True)
        self._save(keep)


# Backwards-compatible name used by earlier callers/tests.
official_images = factory_images