"""Kirin Download-Mode (HiSilicon) — hisi-idt-Protokoll.

Wenn ein Kirin-Gerät per Testpoint zur Erde gezogen und gebootet wird,
landet es im "DLOAD"-Modus mit USB-VID:PID 12d1:1100. In diesem Modus
hat es selbst noch keinen Speicher initialisiert — es wartet darauf,
dass der Host einen *xloader* (kleiner sekundärer Bootloader) per USB
schickt. Erst danach kommt der nächste Stage-Loader (`usb_loader.bin`),
und dann fährt das Gerät als erweiterter Fastboot hoch.

Das Protokoll dafür heißt im Huawei-Kosmos schlicht "hisi-idt". Es ist
nicht öffentlich dokumentiert, aber seit Jahren bekannt aus dem
hisi-idt.py-Skript, das u.a. im Kirin-Tools-Umfeld kursiert. Hier eine
saubere Implementation gegen pyusb.

WICHTIG:
    Die Loader-Dateien (`hisi-sec_usb_xloader.bin`, `usb_loader.bin`)
    sind Huawei-signierte Binaries. Du musst sie selbst beschaffen
    und in `loaders/kirin/<soc>/` ablegen. Siehe loaders/README.md.

Diese Implementation orientiert sich am öffentlichen Protokoll-Wissen.
Auf realer P10-Lite-Hardware vor produktivem Einsatz validieren.
"""
from __future__ import annotations

import struct
import time
from pathlib import Path

import usb.core
import usb.util

KIRIN_DLOAD_VID = 0x12D1
KIRIN_DLOAD_PID = 0x1100

# Protokoll-Konstanten (hisi-idt)
HEAD_TAG = 0xFE
TAIL_TAG = 0xFE
TYPE_DATA = 0xDA
TYPE_HEAD = 0xA5
CHUNK = 0x800  # 2 KiB Payload pro Frame
TIMEOUT_MS = 5000


def _crc16_xmodem(data: bytes) -> int:
    """CRC-16/XMODEM (Polynom 0x1021, init 0)."""
    crc = 0
    for b in data:
        crc ^= b << 8
        for _ in range(8):
            if crc & 0x8000:
                crc = (crc << 1) ^ 0x1021
            else:
                crc <<= 1
            crc &= 0xFFFF
    return crc


def _frame(seq: int, payload: bytes, frame_type: int) -> bytes:
    """Ein hisi-idt-Frame bauen.

    Layout:
        HEAD (1) | TYPE (1) | SEQ (2 BE) | PAYLOAD | CRC16 (2 BE) | TAIL (1)
    """
    body = struct.pack(">BBH", HEAD_TAG, frame_type, seq) + payload
    crc = _crc16_xmodem(body[1:])  # CRC ohne HEAD
    return body + struct.pack(">HB", crc, TAIL_TAG)


class KirinDload:
    """Verbindung zum Kirin im Download-Mode."""

    def __init__(self) -> None:
        self.dev: usb.core.Device | None = None
        self.ep_out = None
        self.ep_in = None

    def open(self) -> None:
        dev = usb.core.find(idVendor=KIRIN_DLOAD_VID, idProduct=KIRIN_DLOAD_PID)
        if dev is None:
            raise RuntimeError(
                f"Kein Kirin-DLOAD-Gerät gefunden "
                f"({KIRIN_DLOAD_VID:04x}:{KIRIN_DLOAD_PID:04x}). "
                f"Testpoint sitzen? USB-Kabel datentauglich?"
            )
        if dev.is_kernel_driver_active(0):
            dev.detach_kernel_driver(0)
        dev.set_configuration()
        cfg = dev.get_active_configuration()
        intf = cfg[(0, 0)]
        self.ep_out = usb.util.find_descriptor(
            intf,
            custom_match=lambda e: usb.util.endpoint_direction(e.bEndpointAddress)
            == usb.util.ENDPOINT_OUT,
        )
        self.ep_in = usb.util.find_descriptor(
            intf,
            custom_match=lambda e: usb.util.endpoint_direction(e.bEndpointAddress)
            == usb.util.ENDPOINT_IN,
        )
        if self.ep_out is None:
            raise RuntimeError("Kein OUT-Endpunkt gefunden — falsche USB-Konfiguration?")
        self.dev = dev

    def close(self) -> None:
        if self.dev is not None:
            usb.util.dispose_resources(self.dev)
            self.dev = None

    def __enter__(self) -> "KirinDload":
        self.open()
        return self

    def __exit__(self, *exc) -> None:
        self.close()

    def send_loader(self, path: Path, load_addr: int) -> None:
        """Loader-Binary in einer Folge von Frames an das Gerät senden.

        Args:
            path: Pfad zur xloader/usb_loader-Datei.
            load_addr: Zieladresse im SRAM. Pro SoC unterschiedlich.
                Für Kirin 658 (P10 Lite, xloader) typischerweise 0x07012000.
                In der Doku des SoCs verifizieren.
        """
        data = path.read_bytes()
        self._send_header(load_addr, len(data))
        seq = 1
        for offset in range(0, len(data), CHUNK):
            chunk = data[offset : offset + CHUNK]
            self._write(_frame(seq, chunk, TYPE_DATA))
            seq = (seq + 1) & 0xFFFF
            # Kleines Delay — Kirin akzeptiert zu schnelle Bursts nicht zuverlässig
            time.sleep(0.001)
        # Sequenz abschließen: einige Implementationen senden ein finales Frame
        # mit Länge 0. Bei Bedarf hier ergänzen, je nach SoC-Verhalten.

    def _send_header(self, addr: int, size: int) -> None:
        """Header-Frame: sagt dem Gerät, wohin und wieviel."""
        payload = struct.pack(">II", addr, size)
        self._write(_frame(0, payload, TYPE_HEAD))

    def _write(self, buf: bytes) -> None:
        assert self.dev is not None and self.ep_out is not None
        self.ep_out.write(buf, timeout=TIMEOUT_MS)