from __future__ import annotations

import plistlib
import re
import subprocess
from collections.abc import Callable, Iterator
from dataclasses import dataclass
from pathlib import PurePosixPath
from typing import TypedDict, cast

from k1link.artifacts import utc_now_iso

USB_IOREG_COMMAND = (
    "/usr/sbin/ioreg",
    "-a",
    "-r",
    "-c",
    "IOUSBHostDevice",
    "-l",
    "-w",
    "0",
)
SERIAL_IOREG_COMMAND = (
    "/usr/sbin/ioreg",
    "-a",
    "-r",
    "-c",
    "IOSerialBSDClient",
    "-l",
    "-w",
    "0",
)
DISKUTIL_COMMAND = ("/usr/sbin/diskutil", "list", "-plist", "external", "physical")


@dataclass(frozen=True)
class CommandOutput:
    argv: tuple[str, ...]
    returncode: int | None
    stdout: bytes
    stderr: bytes
    error: str | None


CommandRunner = Callable[[list[str]], CommandOutput]


class SourceStatus(TypedDict):
    name: str
    argv: list[str]
    ok: bool
    returncode: int | None
    error: str | None


class UsbInterfaceRecord(TypedDict):
    name: str | None
    function_hints: list[str]
    interface_number: int | None
    interface_class: int | None
    interface_class_hex: str | None
    interface_class_name: str | None
    interface_subclass: int | None
    interface_protocol: int | None
    alternate_setting: int | None
    configuration_value: int | None
    endpoint_count: int | None


class UsbDeviceRecord(TypedDict):
    product_name: str | None
    vendor_name: str | None
    serial_number: str | None
    vendor_id: int | None
    vendor_id_hex: str | None
    product_id: int | None
    product_id_hex: str | None
    device_class: int | None
    device_subclass: int | None
    device_protocol: int | None
    usb_bcd: int | None
    device_bcd: int | None
    usb_speed: int | None
    link_speed_bits_per_second: int | None
    usb_address: int | None
    location_id: int | None
    registry_entry_id: int | None
    bsd_names: list[str]
    interface_capabilities: list[str]
    interfaces: list[UsbInterfaceRecord]


class ExternalStorageEntry(TypedDict):
    device_identifier: str
    parent_device_identifier: str | None
    content: str | None
    size_bytes: int | None
    volume_name: str | None
    mount_point: str | None
    os_internal: bool | None
    xgrids_related: bool


class ExternalStorageRecord(TypedDict):
    all_disks: list[str]
    whole_disks: list[str]
    volumes_from_disks: list[str]
    entries: list[ExternalStorageEntry]


class UsbModemRecord(TypedDict):
    callout_device: str | None
    dialin_device: str | None
    tty_base_name: str | None
    client_type: str | None


class SafetyRecord(TypedDict):
    metadata_only: bool
    sudo_used: bool
    device_file_contents_read: bool
    device_writes_performed: bool


class UsbSnapshot(TypedDict):
    schema_version: int
    created_at_utc: str
    sensitivity: str
    safety: SafetyRecord
    sources: list[SourceStatus]
    xgrids_device_count: int
    xgrids_devices: list[UsbDeviceRecord]
    external_storage: ExternalStorageRecord
    usbmodem_device_names: list[str]
    usbmodem_devices: list[UsbModemRecord]


def run_command(argv: list[str]) -> CommandOutput:
    """Run a fixed read-only macOS metadata command without privilege escalation."""
    try:
        result = subprocess.run(
            argv,
            check=False,
            capture_output=True,
            timeout=20,
        )
    except (OSError, subprocess.SubprocessError) as exc:
        return CommandOutput(
            argv=tuple(argv),
            returncode=None,
            stdout=b"",
            stderr=b"",
            error=f"{type(exc).__name__}: {exc}",
        )
    return CommandOutput(
        argv=tuple(argv),
        returncode=result.returncode,
        stdout=result.stdout,
        stderr=result.stderr,
        error=None,
    )


def _node(value: object) -> dict[str, object] | None:
    if not isinstance(value, dict):
        return None
    raw = cast("dict[object, object]", value)
    return {key: item for key, item in raw.items() if isinstance(key, str)}


def _items(value: object) -> list[object]:
    if not isinstance(value, list):
        return []
    return cast("list[object]", value)


def _walk_registry_nodes(value: object) -> Iterator[dict[str, object]]:
    for item in _items(value):
        yield from _walk_registry_nodes(item)

    node = _node(value)
    if node is None:
        return
    yield node
    yield from _walk_registry_nodes(node.get("IORegistryEntryChildren"))


def _string(node: dict[str, object], *keys: str) -> str | None:
    for key in keys:
        value = node.get(key)
        if isinstance(value, str):
            return value
    return None


def _integer(node: dict[str, object], key: str) -> int | None:
    value = node.get(key)
    if isinstance(value, int) and not isinstance(value, bool):
        return value
    return None


def _boolean(node: dict[str, object], key: str) -> bool | None:
    value = node.get(key)
    return value if isinstance(value, bool) else None


def _hex_id(value: int | None) -> str | None:
    return None if value is None else f"0x{value:04x}"


def _is_xgrids_device(node: dict[str, object]) -> bool:
    names = [
        _string(node, "USB Product Name"),
        _string(node, "kUSBProductString"),
        _string(node, "IORegistryEntryName"),
        _string(node, "USB Vendor Name"),
        _string(node, "kUSBVendorString"),
    ]
    identity = " ".join(name.casefold() for name in names if name is not None)
    return "xgrids" in identity or "lixel" in identity


def _usb_class_name(interface_class: int | None) -> str | None:
    if interface_class is None:
        return None
    return {
        0x02: "communications_and_cdc_control",
        0x08: "mass_storage",
        0x0A: "cdc_data",
        0xE0: "wireless_controller",
        0xFF: "vendor_specific",
    }.get(interface_class)


def _function_hints(
    name: str | None,
    interface_class: int | None,
    interface_subclass: int | None,
) -> list[str]:
    normalized = (name or "").casefold()
    hints: set[str] = set()
    if "rndis" in normalized:
        hints.add("rndis")
    if "mass storage" in normalized or interface_class == 0x08:
        hints.add("mass_storage")
    if "ncm" in normalized or (interface_class == 0x02 and interface_subclass == 0x0D):
        hints.add("ncm")
    if any(marker in normalized for marker in ("serial", "modem", "acm")) or (
        interface_class == 0x02 and interface_subclass == 0x02
    ):
        hints.add("serial")
    if interface_class == 0x0A:
        hints.add("cdc_data")
    return sorted(hints)


def _interface_record(node: dict[str, object]) -> UsbInterfaceRecord:
    name = _string(node, "kUSBString", "IORegistryEntryName")
    interface_class = _integer(node, "bInterfaceClass")
    interface_subclass = _integer(node, "bInterfaceSubClass")
    return {
        "name": name,
        "function_hints": _function_hints(name, interface_class, interface_subclass),
        "interface_number": _integer(node, "bInterfaceNumber"),
        "interface_class": interface_class,
        "interface_class_hex": _hex_id(interface_class),
        "interface_class_name": _usb_class_name(interface_class),
        "interface_subclass": interface_subclass,
        "interface_protocol": _integer(node, "bInterfaceProtocol"),
        "alternate_setting": _integer(node, "bAlternateSetting"),
        "configuration_value": _integer(node, "bConfigurationValue"),
        "endpoint_count": _integer(node, "bNumEndpoints"),
    }


def _interface_sort_key(record: UsbInterfaceRecord) -> tuple[bool, int, str]:
    number = record["interface_number"]
    return (number is None, number if number is not None else 0, record["name"] or "")


def _device_record(node: dict[str, object]) -> UsbDeviceRecord:
    interfaces = [
        _interface_record(child)
        for child in _walk_registry_nodes(node.get("IORegistryEntryChildren"))
        if _string(child, "IOObjectClass") == "IOUSBHostInterface"
    ]
    interfaces.sort(key=_interface_sort_key)

    bsd_names = sorted(
        {
            name
            for child in _walk_registry_nodes(node.get("IORegistryEntryChildren"))
            if (name := _string(child, "BSD Name")) is not None
        }
    )
    interface_capabilities = sorted(
        {hint for interface in interfaces for hint in interface["function_hints"]}
    )
    vendor_id = _integer(node, "idVendor")
    product_id = _integer(node, "idProduct")
    return {
        "product_name": _string(
            node, "USB Product Name", "kUSBProductString", "IORegistryEntryName"
        ),
        "vendor_name": _string(node, "USB Vendor Name", "kUSBVendorString"),
        "serial_number": _string(node, "USB Serial Number", "kUSBSerialNumberString"),
        "vendor_id": vendor_id,
        "vendor_id_hex": _hex_id(vendor_id),
        "product_id": product_id,
        "product_id_hex": _hex_id(product_id),
        "device_class": _integer(node, "bDeviceClass"),
        "device_subclass": _integer(node, "bDeviceSubClass"),
        "device_protocol": _integer(node, "bDeviceProtocol"),
        "usb_bcd": _integer(node, "bcdUSB"),
        "device_bcd": _integer(node, "bcdDevice"),
        "usb_speed": _integer(node, "USBSpeed"),
        "link_speed_bits_per_second": _integer(node, "UsbLinkSpeed"),
        "usb_address": _integer(node, "USB Address"),
        "location_id": _integer(node, "locationID"),
        "registry_entry_id": _integer(node, "IORegistryEntryID"),
        "bsd_names": bsd_names,
        "interface_capabilities": interface_capabilities,
        "interfaces": interfaces,
    }


def parse_xgrids_devices(plist: object) -> list[UsbDeviceRecord]:
    devices = [
        _device_record(node)
        for node in (_node(item) for item in _items(plist))
        if node is not None and _is_xgrids_device(node)
    ]
    devices.sort(
        key=lambda record: (
            record["product_name"] or "",
            record["serial_number"] or "",
            record["location_id"] or 0,
        )
    )
    return devices


def _string_list(node: dict[str, object], key: str) -> list[str]:
    return sorted(item for item in _items(node.get(key)) if isinstance(item, str))


def _disk_root(device_identifier: str) -> str | None:
    match = re.fullmatch(r"(disk\d+)(?:s\d+)*", device_identifier)
    return None if match is None else match.group(1)


def _external_storage_entries(
    value: object,
    parent_device_identifier: str | None,
    xgrids_disk_roots: set[str],
) -> list[ExternalStorageEntry]:
    entries: list[ExternalStorageEntry] = []
    for item in _items(value):
        node = _node(item)
        if node is None:
            continue
        device_identifier = _string(node, "DeviceIdentifier")
        next_parent = parent_device_identifier
        if device_identifier is not None:
            disk_root = _disk_root(device_identifier)
            entries.append(
                {
                    "device_identifier": device_identifier,
                    "parent_device_identifier": parent_device_identifier,
                    "content": _string(node, "Content"),
                    "size_bytes": _integer(node, "Size"),
                    "volume_name": _string(node, "VolumeName"),
                    "mount_point": _string(node, "MountPoint"),
                    "os_internal": _boolean(node, "OSInternal"),
                    "xgrids_related": bool(
                        disk_root is not None and disk_root in xgrids_disk_roots
                    ),
                }
            )
            next_parent = device_identifier
        for child_key in ("Partitions", "APFSVolumes"):
            entries.extend(
                _external_storage_entries(
                    node.get(child_key),
                    next_parent,
                    xgrids_disk_roots,
                )
            )
    return entries


def parse_external_storage(
    plist: object,
    xgrids_devices: list[UsbDeviceRecord],
) -> ExternalStorageRecord:
    node = _node(plist) or {}
    xgrids_disk_roots = {
        root
        for device in xgrids_devices
        for name in device["bsd_names"]
        if (root := _disk_root(name)) is not None
    }
    entries = _external_storage_entries(
        node.get("AllDisksAndPartitions"),
        None,
        xgrids_disk_roots,
    )
    entries.sort(key=lambda entry: entry["device_identifier"])
    return {
        "all_disks": _string_list(node, "AllDisks"),
        "whole_disks": _string_list(node, "WholeDisks"),
        "volumes_from_disks": _string_list(node, "VolumesFromDisks"),
        "entries": entries,
    }


def _is_usbmodem_path(path: str | None) -> bool:
    return path is not None and PurePosixPath(path).name.startswith(("cu.usbmodem", "tty.usbmodem"))


def parse_usbmodem_devices(plist: object) -> tuple[list[str], list[UsbModemRecord]]:
    names: set[str] = set()
    records: list[UsbModemRecord] = []
    for item in _items(plist):
        node = _node(item)
        if node is None:
            continue
        callout = _string(node, "IOCalloutDevice")
        dialin = _string(node, "IODialinDevice")
        matched_paths: list[str] = []
        for path in (callout, dialin):
            if path is not None and _is_usbmodem_path(path):
                matched_paths.append(path)
        if not matched_paths:
            continue
        names.update(matched_paths)
        records.append(
            {
                "callout_device": callout,
                "dialin_device": dialin,
                "tty_base_name": _string(node, "IOTTYBaseName"),
                "client_type": _string(node, "IOSerialBSDClientType"),
            }
        )
    records.sort(key=lambda record: (record["callout_device"] or "", record["dialin_device"] or ""))
    return sorted(names), records


def _source_status(name: str, output: CommandOutput) -> SourceStatus:
    error = output.error
    if error is None and output.returncode != 0:
        stderr = output.stderr.decode("utf-8", errors="replace").strip()
        error = stderr[:500] or f"command exited with status {output.returncode}"
    return {
        "name": name,
        "argv": list(output.argv),
        "ok": error is None and output.returncode == 0,
        "returncode": output.returncode,
        "error": error,
    }


def _load_plist(
    name: str,
    output: CommandOutput,
    *,
    empty_is_list: bool = False,
) -> tuple[object | None, SourceStatus]:
    status = _source_status(name, output)
    if not status["ok"]:
        return None, status
    if empty_is_list and not output.stdout.strip():
        return [], status
    try:
        return plistlib.loads(output.stdout), status
    except (plistlib.InvalidFileException, ValueError, TypeError, OverflowError) as exc:
        status["ok"] = False
        status["error"] = f"invalid plist: {type(exc).__name__}: {exc}"
        return None, status


def snapshot(runner: CommandRunner = run_command) -> UsbSnapshot:
    """Collect USB registry and storage metadata without opening device files."""
    usb_output = runner(list(USB_IOREG_COMMAND))
    serial_output = runner(list(SERIAL_IOREG_COMMAND))
    storage_output = runner(list(DISKUTIL_COMMAND))

    usb_plist, usb_status = _load_plist("usb_ioreg", usb_output, empty_is_list=True)
    serial_plist, serial_status = _load_plist("serial_ioreg", serial_output, empty_is_list=True)
    storage_plist, storage_status = _load_plist("external_disks", storage_output)

    devices = parse_xgrids_devices(usb_plist)
    usbmodem_names, usbmodem_devices = parse_usbmodem_devices(serial_plist)
    external_storage = parse_external_storage(storage_plist, devices)
    return {
        "schema_version": 1,
        "created_at_utc": utc_now_iso(),
        "sensitivity": (
            "contains USB serial numbers, BSD device names and external volume metadata; "
            "store only in an ignored session path and do not commit"
        ),
        "safety": {
            "metadata_only": True,
            "sudo_used": False,
            "device_file_contents_read": False,
            "device_writes_performed": False,
        },
        "sources": [usb_status, serial_status, storage_status],
        "xgrids_device_count": len(devices),
        "xgrids_devices": devices,
        "external_storage": external_storage,
        "usbmodem_device_names": usbmodem_names,
        "usbmodem_devices": usbmodem_devices,
    }
