from __future__ import annotations

import hashlib
import hmac
import tarfile
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import IO

from k1link.host_network.wifi import (
    HostWifiCredentialMaterialStoreResult,
    store_wifi_credential_material,
)

K1_FW302_CREDENTIAL_PROVIDER_ID = (
    "xgrids.lixelkity-k1.quick-connect.fw-3.0.2.official-firmware.v1"
)
K1_FW302_OFFICIAL_ARCHIVE_SHA256 = (
    "e5830feae54d586cdeda2824495d08598920dc9cf4541059d01f0efeb858a750"
)
K1_FW302_OUTER_MEMBER = "upgrade.tar.gz"
K1_FW302_RK_IMAGE_MEMBER = (
    "upgrade/rk/rk_normal/"
    "OTA-BOOT-ROOTFS-APP-V3.0.2-20250624.153447.img"
)
K1_FW302_APPS_OFFSET = 4_338_829_862
K1_FW302_APPS_SIZE = 230_801_408

_AP_PSK_DECLARATION = b"nmcli conn modify WIFI_AP 802-11-wireless-security.psk "
_READ_CHUNK_SIZE = 4 * 1024 * 1024
_SCAN_OVERLAP = len(_AP_PSK_DECLARATION) + 128


class FirmwareCredentialError(RuntimeError):
    """An exact firmware artifact cannot provide the reviewed credential."""


class SecretBuffer:
    """Short-lived credential bytes whose repr and str are always redacted."""

    __slots__ = ("_value",)

    def __init__(self, value: bytearray) -> None:
        self._value = value

    def __repr__(self) -> str:
        return "SecretBuffer(<redacted>)"

    __str__ = __repr__

    def reveal_ascii(self) -> str:
        return self._value.decode("ascii")

    def zeroize(self) -> None:
        self._value[:] = b"\x00" * len(self._value)


@dataclass(frozen=True)
class FirmwareCredentialImportResult:
    provider_id: str
    firmware_sha256: str
    host_adapter: str
    outcome: str


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        while chunk := stream.read(_READ_CHUNK_SIZE):
            digest.update(chunk)
    return digest.hexdigest()


def _discard_exact(stream: IO[bytes], size: int) -> None:
    remaining = size
    while remaining:
        chunk = stream.read(min(_READ_CHUNK_SIZE, remaining))
        if not chunk:
            raise FirmwareCredentialError("firmware image ended before the apps partition")
        remaining -= len(chunk)


def _bounded_chunks(stream: IO[bytes], size: int) -> Iterator[bytes]:
    remaining = size
    while remaining:
        chunk = stream.read(min(_READ_CHUNK_SIZE, remaining))
        if not chunk:
            raise FirmwareCredentialError("firmware apps partition is truncated")
        remaining -= len(chunk)
        yield chunk


def _credential_from_chunks(chunks: Iterator[bytes]) -> SecretBuffer:
    overlap = b""
    processed = 0
    matches: dict[int, bytearray] = {}
    for chunk in chunks:
        combined = overlap + chunk
        combined_offset = processed - len(overlap)
        start = 0
        while True:
            index = combined.find(_AP_PSK_DECLARATION, start)
            if index < 0:
                break
            value_start = index + len(_AP_PSK_DECLARATION)
            value_end = value_start
            while value_end < len(combined) and combined[value_end] not in b"\x00\r\n":
                value_end += 1
            if value_end < len(combined):
                absolute_offset = combined_offset + index
                matches[absolute_offset] = bytearray(combined[value_start:value_end].strip())
            start = index + 1
        processed += len(chunk)
        overlap = combined[-_SCAN_OVERLAP:]

    if len(matches) != 1:
        for value in matches.values():
            value[:] = b"\x00" * len(value)
        raise FirmwareCredentialError(
            "reviewed AP credential declaration was not unique in the apps partition"
        )
    value = next(iter(matches.values()))
    if not 8 <= len(value) <= 63 or any(byte <= 0x20 or byte >= 0x7F for byte in value):
        value[:] = b"\x00" * len(value)
        raise FirmwareCredentialError("reviewed AP credential has an invalid WPA-PSK shape")
    return SecretBuffer(value)


@contextmanager
def _official_rk_image(archive_path: Path) -> Iterator[IO[bytes]]:
    try:
        with tarfile.open(archive_path, mode="r:*") as outer:
            outer_member = outer.getmember(K1_FW302_OUTER_MEMBER)
            upgrade_stream = outer.extractfile(outer_member)
            if upgrade_stream is None:
                raise FirmwareCredentialError("official upgrade member has no readable payload")
            with upgrade_stream, tarfile.open(
                fileobj=upgrade_stream, mode="r|gz"
            ) as upgrade:
                for member in upgrade:
                    if member.name != K1_FW302_RK_IMAGE_MEMBER:
                        continue
                    image_stream = upgrade.extractfile(member)
                    if image_stream is None:
                        raise FirmwareCredentialError(
                            "official Rockchip image has no readable payload"
                        )
                    with image_stream:
                        yield image_stream
                    return
    except (KeyError, OSError, tarfile.TarError) as exc:
        raise FirmwareCredentialError("official firmware archive is unreadable") from exc
    raise FirmwareCredentialError("reviewed Rockchip image is absent from the archive")


def extract_k1_fw302_ap_credential(archive_path: Path) -> tuple[SecretBuffer, str]:
    """Resolve the FW 3.0.2 AP material without printing or persisting it.

    The exact official archive is authenticated first. Only then is the reviewed
    apps-partition range scanned for one bounded NetworkManager declaration.
    """

    resolved_path = archive_path.resolve(strict=True)
    firmware_sha256 = _sha256(resolved_path)
    if not hmac.compare_digest(firmware_sha256, K1_FW302_OFFICIAL_ARCHIVE_SHA256):
        raise FirmwareCredentialError("firmware SHA-256 does not match the reviewed 3.0.2 image")
    with _official_rk_image(resolved_path) as image_stream:
        _discard_exact(image_stream, K1_FW302_APPS_OFFSET)
        secret = _credential_from_chunks(
            _bounded_chunks(image_stream, K1_FW302_APPS_SIZE)
        )
    return secret, firmware_sha256


def import_k1_fw302_ap_credential(
    archive_path: Path,
    helper_path: Path,
) -> FirmwareCredentialImportResult:
    """Install one firmware-scoped material in the host's secure store."""

    secret, firmware_sha256 = extract_k1_fw302_ap_credential(archive_path)
    try:
        stored: HostWifiCredentialMaterialStoreResult = store_wifi_credential_material(
            helper_path,
            K1_FW302_CREDENTIAL_PROVIDER_ID,
            secret.reveal_ascii(),
        )
    finally:
        secret.zeroize()
    return FirmwareCredentialImportResult(
        provider_id=K1_FW302_CREDENTIAL_PROVIDER_ID,
        firmware_sha256=firmware_sha256,
        host_adapter=stored["adapter"],
        outcome=stored["outcome"],
    )
