"""Firmware images: validation, integrity, layout and the firmware folder.

The tool flashes ONLY the .bin files in its own firmware/ folder (next to
kota_tech/). There is no import and no per-user library: to offer a build,
the firmware team puts the file into that folder. Images are listed by the
compile date stored in the firmware itself (ESP-IDF app descriptor), newest
first, so the file name does not matter.

Two kinds:
  * merged - bootloader at 0x0 + partition table + app (merged-binary.bin);
             App-only or Full flashing;
  * app    - the application binary only (build/kota-sensorboard.bin);
             App-only flashing, laid out like a merged image in the folder.

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 struct
from dataclasses import dataclass
from pathlib import Path

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

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          # "Sep 28 2026 15:07:57" 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

    @property
    def compiled(self) -> str:
        """'2026-09-28 15:07' (ISO-like, sortable) or the raw descriptor text."""
        t = self.built_at
        return t.strftime("%Y-%m-%d %H:%M") if t else self.built


@dataclass(frozen=True)
class FirmwareImage:
    path: Path
    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 name(self) -> str:
        return self.path.name

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

    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:
        """Choice text: compile date first (that is what the technician picks by)."""
        kind = "" if self.kind == "merged" else ", app only"
        return f"Compiled {self.app.compiled}  -  v{self.app.version} [{self.app.elf_sha[:8]}]{kind}  ({self.name})"


# ---------------------------------------------------------------- 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, reference: FirmwareImage | None = None) -> FirmwareImage:
    """Validate a merged or app-only image. App-only images need a `reference`
    (a merged 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 file, but there is no complete (merged) image in the folder "
                             "to take the flash layout from")
        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, kind, hashlib.sha256(data).hexdigest(),
                         flash_size, parts, table, app_off, length, app)


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

def folder_images(folder: Path = FIRMWARE_DIR) -> tuple[list[FirmwareImage], list[str]]:
    """(valid images newest-compiled first, problems) for the .bin files in `folder`.

    A broken or foreign file is reported, never offered. The same build
    under two file names is offered once.
    """
    if not folder.is_dir():
        return [], [f"firmware folder not found: {folder}"]
    files = sorted(p for p in folder.iterdir() if p.is_file() and p.suffix.lower() == ".bin")
    merged: list[FirmwareImage] = []
    app_only: list[Path] = []
    problems: list[str] = []
    for f in files:
        try:
            merged.append(inspect_image(f))
        except ImageError as exc:
            if "app-only" in str(exc):
                app_only.append(f)
            else:
                problems.append(f"{f.name}: {exc}")
        except OSError as exc:
            problems.append(f"{f.name}: cannot read ({exc})")
    merged.sort(key=_newest_first)
    reference = merged[0] if merged else None
    images = list(merged)
    for f in app_only:
        try:
            images.append(inspect_image(f, reference))
        except (ImageError, OSError) as exc:
            problems.append(f"{f.name}: {exc}")
    images.sort(key=_newest_first)
    unique: dict[str, FirmwareImage] = {}
    for img in images:
        if img.sha256 in unique:
            problems.append(f"{img.name}: same file as {unique[img.sha256].name} (listed once)")
        else:
            unique[img.sha256] = img
    return list(unique.values()), problems


def _newest_first(img: FirmwareImage) -> tuple:
    # Unknown compile dates last; merged before app-only for the same date.
    t = img.app.built_at
    return (t is None, -t.timestamp() if t else 0.0, img.kind != "merged", img.name)