"""Firmware images for ninlight lamps: trailer, version order, Wi-Fi check, catalog.

Deliberately free of Home Assistant imports. Blocking functions are marked; callers run them in an
executor. Facts behind the constants: spec §2 (layout checked against a real image).
"""

from __future__ import annotations

from dataclasses import dataclass
import hashlib
import logging
from pathlib import Path
import re
import zlib

_LOGGER = logging.getLogger(__name__)

TRAILER_SIZE = 512
TRAILER_MAGIC = b"NINL"
TRAILER_FORMAT = 1
W600_HEADER_SIZE = 56
GZIP_MAGIC = b"\x1f\x8b\x08"
# cmake/write_version.cmake cuts the version hard to 15 characters.
VERSION_FIELD_MAX = 15

_VERSION_RE = re.compile(
    r"^(?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)"
    r"(?:-rc(?P<rc>\d+))?"
    r"(?:-(?P<commits>\d+)(?P<hash>-g[0-9a-f]*)?)?"
    r"(?:-d(?:i(?:r(?:t(?:y)?)?)?)?)?$"
)


class InvalidImageError(ValueError):
    """The file is not a valid ninlight OTA image."""


@dataclass(frozen=True, slots=True)
class FirmwareVersion:
    """A parsed `git describe` version. The hash is ignored on purpose (usually truncated)."""

    base: tuple[int, int, int]
    rc: int | None
    commits: int

    @property
    def sort_key(self) -> tuple[int, int, int, int, int, int]:
        # A release candidate sorts before the release of the same base.
        rc_rank, rc_number = (0, self.rc) if self.rc is not None else (1, 0)
        return (*self.base, rc_rank, rc_number, self.commits)


def parse_version(text: str) -> FirmwareVersion | None:
    match = _VERSION_RE.match(text)
    if match is None:
        return None
    if len(text) == VERSION_FIELD_MAX and text[-1].isdigit() and match["hash"] is None:
        # The cut may have landed inside the commit (or rc) number: not comparable.
        return None
    return FirmwareVersion(
        (int(match["major"]), int(match["minor"]), int(match["patch"])),
        int(match["rc"]) if match["rc"] is not None else None,
        int(match["commits"]) if match["commits"] is not None else 0,
    )


def is_newer(candidate: str, installed: str) -> bool:
    new, old = parse_version(candidate), parse_version(installed)
    if new is None or old is None:
        return False
    return new.sort_key > old.sort_key


def parse_trailer(data: bytes) -> tuple[str, str]:
    """Return (device_type, version) of a valid image or raise InvalidImageError."""
    if len(data) <= TRAILER_SIZE:
        raise InvalidImageError("too short for a trailer")
    payload, trailer = data[:-TRAILER_SIZE], data[-TRAILER_SIZE:]
    if trailer[0:4] != TRAILER_MAGIC:
        raise InvalidImageError("no NINL trailer")
    if trailer[4] != TRAILER_FORMAT:
        raise InvalidImageError(f"unsupported trailer format {trailer[4]}")
    if trailer[37:69] != hashlib.sha256(payload).digest():
        raise InvalidImageError("sha256 mismatch")
    device_type = _text_field(trailer[5:21])
    version = _text_field(trailer[21:37])
    if not device_type or not version:
        raise InvalidImageError("empty device type or version")
    return device_type, version


def _text_field(raw: bytes) -> str:
    return raw.split(b"\x00", 1)[0].decode("ascii", errors="replace")


def contains_ssid(data: bytes, ssid: str) -> bool:
    """True if the decompressed firmware carries `ssid` as a NUL-framed C string (spec §7)."""
    if not ssid or len(data) <= TRAILER_SIZE:
        return False
    payload = data[:-TRAILER_SIZE]
    start = payload.find(GZIP_MAGIC, W600_HEADER_SIZE)
    if start < 0:
        return False
    try:
        raw = zlib.decompressobj(31).decompress(payload[start:])
    except zlib.error:
        return False
    return b"\x00" + ssid.encode() + b"\x00" in raw


@dataclass(frozen=True, slots=True)
class FirmwareImage:
    """One valid image in the firmware folder."""

    path: Path
    device_type: str
    version: str
    size: int
    mtime_ns: int


def _sort_key(image: FirmwareImage) -> tuple[int, ...]:
    version = parse_version(image.version)
    return version.sort_key if version is not None else ()


class FirmwareCatalog:
    """All valid images in one folder. Not thread-safe: callers serialize scans."""

    def __init__(self, directory: Path) -> None:
        self.directory = directory
        self.images: tuple[FirmwareImage, ...] = ()
        self._parsed: dict[tuple[str, int, int], FirmwareImage | None] = {}
        self._ssid_hits: dict[tuple[str, int, int, str], bool] = {}

    def scan(self) -> tuple[FirmwareImage, ...]:
        """Blocking. Invalid files are logged once per (path, mtime, size)."""
        self.directory.mkdir(parents=True, exist_ok=True)
        found: list[FirmwareImage] = []
        seen: set[tuple[str, int, int]] = set()
        for path in sorted(self.directory.glob("*.img")):
            try:
                if not path.is_file():
                    continue
                stat = path.stat()
            except OSError:
                continue
            key = (str(path), stat.st_mtime_ns, stat.st_size)
            seen.add(key)
            if key not in self._parsed:
                try:
                    device_type, version = parse_trailer(path.read_bytes())
                except (OSError, InvalidImageError) as err:
                    _LOGGER.warning("Ignoring firmware file %s: %s", path.name, err)
                    self._parsed[key] = None
                else:
                    self._parsed[key] = FirmwareImage(
                        path, device_type, version, stat.st_size, stat.st_mtime_ns
                    )
            image = self._parsed[key]
            if image is not None:
                found.append(image)
        self._parsed = {k: v for k, v in self._parsed.items() if k in seen}
        self._ssid_hits = {k: v for k, v in self._ssid_hits.items() if k[:3] in seen}
        self.images = tuple(found)
        return self.images

    def newer_images(self, device_type: str, installed: str) -> list[FirmwareImage]:
        """Images for exactly this device type that are newer than `installed`, newest first."""
        candidates = [
            image
            for image in self.images
            if image.device_type == device_type and is_newer(image.version, installed)
        ]
        return sorted(candidates, key=_sort_key, reverse=True)

    def verify(self, image: FirmwareImage) -> bool:
        """Blocking. The file is still there, unchanged and valid for the same type/version."""
        try:
            stat = image.path.stat()
            if (stat.st_mtime_ns, stat.st_size) != (image.mtime_ns, image.size):
                return False
            return parse_trailer(image.path.read_bytes()) == (image.device_type, image.version)
        except OSError, InvalidImageError:
            return False

    def contains_ssid(self, image: FirmwareImage, ssid: str) -> bool:
        """Blocking, cached per (image file, ssid)."""
        key = (str(image.path), image.mtime_ns, image.size, ssid)
        if key not in self._ssid_hits:
            try:
                data = image.path.read_bytes()
            except OSError:
                return False
            self._ssid_hits[key] = contains_ssid(data, ssid)
        return self._ssid_hits[key]