from __future__ import annotations

from contextlib import contextmanager
from dataclasses import dataclass
import mmap
import os
from pathlib import Path
import re
import struct
import time

from .firmware import FirmwareImage


ASMEDIA_VENDOR_ID = 0x1B21
XHCI_CLASS_CODE = 0x0C0330
KNOWN_DEVICE_IDS = {0x3042, 0x3142}
KNOWN_PROTOCOL_DEVICE_IDS = {0x3042, 0x3142}
ROM_FW_VERSION = 0x010250090816

# PCI configuration space
ASMT_CFG_CONTROL = 0xE0
ASMT_CFG_CONTROL_WRITE = 1 << 1
ASMT_CFG_CONTROL_READ = 1 << 0
ASMT_CFG_SRAM_ADDR = 0xE2
ASMT_CFG_SRAM_ACCESS = 0xEF
ASMT_CFG_SRAM_ACCESS_READ = 1 << 6
ASMT_CFG_SRAM_ACCESS_ENABLE = 1 << 7
ASMT_CFG_DATA_READ0 = 0xF0
ASMT_CFG_DATA_READ1 = 0xF4
ASMT_CFG_DATA_WRITE0 = 0xF8
ASMT_CFG_DATA_WRITE1 = 0xFC

# Mailbox / firmware commands
ASMT_CMD_GET_FWVER = 0x8000060840

# BAR0 registers
ASMT_REG_ADDR = 0x3000
ASMT_REG_WDATA = 0x3004
ASMT_REG_RDATA = 0x3008
ASMT_REG_STATUS = 0x3009
ASMT_REG_STATUS_BUSY = 1 << 7
ASMT_REG_CODE_WDATA = 0x3010

ASMT_MMIO_CPU_MISC = 0x500E
ASMT_MMIO_CPU_MISC_CODE_RAM_WR = 1 << 0
ASMT_MMIO_CPU_MODE_NEXT = 0x5040
ASMT_MMIO_CPU_MODE_RAM = 1 << 0
ASMT_MMIO_CPU_MODE_HALFSPEED = 1 << 1
ASMT_MMIO_CPU_EXEC_CTRL = 0x5042
ASMT_MMIO_CPU_EXEC_CTRL_RESET = 1 << 0
ASMT_MMIO_CPU_EXEC_CTRL_HALT = 1 << 1

XHCI_USBCMD_HCRST = 1 << 1

MAILBOX_TIMEOUT_SEC = 0.010
RESET_TIMEOUT_SEC = 0.500
REG_TIMEOUT_SEC = 0.010
POLL_INTERVAL_SEC = 0.0005

PCI_BDF_RE = re.compile(
    r"^(?:(?P<domain>[0-9a-fA-F]{4}):)?(?P<bus>[0-9a-fA-F]{2}):(?P<slot>[0-9a-fA-F]{2})\.(?P<func>[0-7])$"
)


class FlasherError(RuntimeError):
    """Base error for package failures."""


class DeviceNotFoundError(FlasherError):
    """Raised when no suitable PCI device is available."""


class DeviceCompatibilityError(FlasherError):
    """Raised when the selected device is not known-safe for the protocol."""


@dataclass(frozen=True)
class DeviceInfo:
    bdf: str
    vendor_id: int
    device_id: int
    class_code: int
    driver: str | None
    sysfs_path: Path

    @property
    def known_device(self) -> bool:
        return self.device_id in KNOWN_DEVICE_IDS

    @property
    def protocol_supported(self) -> bool:
        return self.device_id in KNOWN_PROTOCOL_DEVICE_IDS


def normalize_bdf(value: str) -> str:
    match = PCI_BDF_RE.fullmatch(value)
    if not match:
        raise ValueError(f"invalid PCI BDF '{value}'")
    domain = match.group("domain") or "0000"
    return (
        f"{domain.lower()}:{match.group('bus').lower()}:"
        f"{match.group('slot').lower()}.{match.group('func')}"
    )


def discover_devices(sysfs_root: str | Path = "/sys/bus/pci/devices") -> list[DeviceInfo]:
    root = Path(sysfs_root)
    devices: list[DeviceInfo] = []

    for path in sorted(root.iterdir()):
        try:
            vendor_id = _read_hex_file(path / "vendor")
            device_id = _read_hex_file(path / "device")
            class_code = _read_hex_file(path / "class")
        except (OSError, ValueError):
            continue

        if vendor_id != ASMEDIA_VENDOR_ID or class_code != XHCI_CLASS_CODE:
            continue

        devices.append(
            DeviceInfo(
                bdf=path.name,
                vendor_id=vendor_id,
                device_id=device_id,
                class_code=class_code,
                driver=_resolve_driver_name(path),
                sysfs_path=path,
            )
        )

    return devices


def pick_device(
    requested_bdf: str | None,
    *,
    sysfs_root: str | Path = "/sys/bus/pci/devices",
) -> DeviceInfo:
    devices = discover_devices(sysfs_root=sysfs_root)

    if requested_bdf:
        normalized = normalize_bdf(requested_bdf)
        for device in devices:
            if device.bdf == normalized:
                return device
        raise DeviceNotFoundError(f"ASMedia xHCI device {normalized} was not found")

    if not devices:
        raise DeviceNotFoundError("no ASMedia xHCI devices were found in /sys/bus/pci/devices")
    if len(devices) > 1:
        found = ", ".join(device.bdf for device in devices)
        raise DeviceNotFoundError(
            f"multiple ASMedia xHCI devices found ({found}); pass --bdf explicitly"
        )
    return devices[0]


class AsmediaXhciController:
    def __init__(self, info: DeviceInfo):
        self.info = info
        self._config_fd: int | None = None
        self._resource_fd: int | None = None
        self._devmem_fd: int | None = None
        self._bar0: mmap.mmap | None = None
        self._bar0_size: int | None = None
        self._bar0_offset: int = 0

    def open(self) -> "AsmediaXhciController":
        self._ensure_enabled()
        try:
            self._config_fd = os.open(self.info.sysfs_path / "config", os.O_RDWR | os.O_SYNC)
            return self
        except FlasherError:
            self.close()
            raise
        except OSError as exc:
            self.close()
            raise FlasherError(f"failed to open PCI config on {self.info.bdf}: {exc}") from exc

    def close(self) -> None:
        if self._bar0 is not None:
            self._bar0.close()
            self._bar0 = None
        self._bar0_offset = 0
        if self._devmem_fd is not None:
            os.close(self._devmem_fd)
            self._devmem_fd = None
        if self._resource_fd is not None:
            os.close(self._resource_fd)
            self._resource_fd = None
        if self._config_fd is not None:
            os.close(self._config_fd)
            self._config_fd = None

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

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

    @contextmanager
    def temporarily_unbound(self, enabled: bool = True):
        driver = self.info.driver if enabled else None
        if driver:
            self._unbind_driver(driver)
        try:
            yield
        finally:
            if driver:
                self._bind_driver(driver)

    def get_firmware_version(self) -> int:
        self._mailbox_tx(ASMT_CMD_GET_FWVER)
        self._mailbox_tx(0)
        reply_command = self._mailbox_rx()
        version = self._mailbox_rx()
        if reply_command != ASMT_CMD_GET_FWVER:
            raise FlasherError(
                f"unexpected mailbox reply 0x{reply_command:010x}, expected 0x{ASMT_CMD_GET_FWVER:010x}"
            )
        return version

    def upload_firmware(self, firmware: FirmwareImage) -> None:
        self._write_reg(ASMT_MMIO_CPU_MODE_NEXT, ASMT_MMIO_CPU_MODE_HALFSPEED, verify=False)
        self._write_reg(ASMT_MMIO_CPU_EXEC_CTRL, ASMT_MMIO_CPU_EXEC_CTRL_RESET, verify=False)
        self._wait_reset()

        self._write_reg(ASMT_MMIO_CPU_EXEC_CTRL, ASMT_MMIO_CPU_EXEC_CTRL_HALT, verify=False)
        self._write_reg(ASMT_MMIO_CPU_MISC, ASMT_MMIO_CPU_MISC_CODE_RAM_WR, verify=True)

        self.write_config_u8(ASMT_CFG_SRAM_ACCESS, ASMT_CFG_SRAM_ACCESS_ENABLE)
        try:
            for addr, data in firmware.iter_sram_writes():
                self.write_config_u16(ASMT_CFG_SRAM_ADDR, addr)
                self.write_mmio_u32(ASMT_REG_CODE_WDATA, data)
                self._poll_until(
                    lambda current_addr=addr: self.read_config_u16(ASMT_CFG_SRAM_ADDR) != current_addr,
                    timeout=REG_TIMEOUT_SEC,
                    message=f"timed out waiting for SRAM write at 0x{addr:04x}",
                )
        finally:
            self.write_config_u8(ASMT_CFG_SRAM_ACCESS, 0)

        self._write_reg(ASMT_MMIO_CPU_MISC, 0, verify=True)
        self._write_reg(
            ASMT_MMIO_CPU_MODE_NEXT,
            ASMT_MMIO_CPU_MODE_RAM | ASMT_MMIO_CPU_MODE_HALFSPEED,
            verify=False,
        )
        self._write_reg(ASMT_MMIO_CPU_EXEC_CTRL, 0, verify=False)
        self._wait_reset()

    def read_config_u8(self, offset: int) -> int:
        return self._pread_exact(offset, 1)[0]

    def read_config_u16(self, offset: int) -> int:
        return struct.unpack("<H", self._pread_exact(offset, 2))[0]

    def read_config_u32(self, offset: int) -> int:
        return struct.unpack("<I", self._pread_exact(offset, 4))[0]

    def write_config_u8(self, offset: int, value: int) -> None:
        self._pwrite_exact(offset, bytes((value & 0xFF,)))

    def write_config_u16(self, offset: int, value: int) -> None:
        self._pwrite_exact(offset, struct.pack("<H", value & 0xFFFF))

    def write_config_u32(self, offset: int, value: int) -> None:
        self._pwrite_exact(offset, struct.pack("<I", value & 0xFFFFFFFF))

    def read_mmio_u8(self, offset: int) -> int:
        self._ensure_bar0_mapped()
        return self._bar_slice(offset, 1)[0]

    def read_mmio_u32(self, offset: int) -> int:
        self._ensure_bar0_mapped()
        return struct.unpack("<I", self._bar_slice(offset, 4))[0]

    def write_mmio_u8(self, offset: int, value: int) -> None:
        self._ensure_bar0_mapped()
        self._bar_write(offset, bytes((value & 0xFF,)))

    def write_mmio_u16(self, offset: int, value: int) -> None:
        self._ensure_bar0_mapped()
        self._bar_write(offset, struct.pack("<H", value & 0xFFFF))

    def write_mmio_u32(self, offset: int, value: int) -> None:
        self._ensure_bar0_mapped()
        self._bar_write(offset, struct.pack("<I", value & 0xFFFFFFFF))

    def _mailbox_tx(self, data: int) -> None:
        self._poll_until(
            lambda: not (self.read_config_u8(ASMT_CFG_CONTROL) & ASMT_CFG_CONTROL_WRITE),
            timeout=MAILBOX_TIMEOUT_SEC,
            message=f"timed out waiting for mailbox TX slot on {self.info.bdf}",
        )
        self.write_config_u32(ASMT_CFG_DATA_WRITE0, data & 0xFFFFFFFF)
        self.write_config_u32(ASMT_CFG_DATA_WRITE1, (data >> 32) & 0xFFFFFFFF)
        self.write_config_u8(ASMT_CFG_CONTROL, ASMT_CFG_CONTROL_WRITE)

    def _mailbox_rx(self) -> int:
        self._poll_until(
            lambda: bool(self.read_config_u8(ASMT_CFG_CONTROL) & ASMT_CFG_CONTROL_READ),
            timeout=MAILBOX_TIMEOUT_SEC,
            message=f"timed out waiting for mailbox RX data on {self.info.bdf}",
        )
        low = self.read_config_u32(ASMT_CFG_DATA_READ0)
        high = self.read_config_u32(ASMT_CFG_DATA_READ1)
        self.write_config_u8(ASMT_CFG_CONTROL, ASMT_CFG_CONTROL_READ)
        return (high << 32) | low

    def _wait_reset(self) -> None:
        command_offset = self.read_mmio_u8(0x00)

        def reset_done() -> bool:
            return not (self.read_mmio_u32(command_offset) & XHCI_USBCMD_HCRST)

        try:
            self._poll_until(
                reset_done,
                timeout=RESET_TIMEOUT_SEC,
                interval=0.001,
                message=f"xHCI reset timed out on {self.info.bdf}",
            )
            return
        except TimeoutError:
            pass

        self.write_config_u8(ASMT_CFG_SRAM_ACCESS, ASMT_CFG_SRAM_ACCESS_ENABLE)
        self.write_config_u8(ASMT_CFG_SRAM_ACCESS, 0)
        self._poll_until(
            reset_done,
            timeout=RESET_TIMEOUT_SEC,
            interval=0.001,
            message=f"xHCI reset timed out on {self.info.bdf} after SRAM access kick",
        )

    def _read_reg(self, addr: int) -> int:
        self._poll_until(
            lambda: not (self.read_mmio_u8(ASMT_REG_STATUS) & ASMT_REG_STATUS_BUSY),
            timeout=REG_TIMEOUT_SEC,
            message=f"timed out waiting for register bus idle before read [0x{addr:04x}]",
        )
        self.write_mmio_u16(ASMT_REG_ADDR, addr)
        self._poll_until(
            lambda: not (self.read_mmio_u8(ASMT_REG_STATUS) & ASMT_REG_STATUS_BUSY),
            timeout=REG_TIMEOUT_SEC,
            message=f"timed out waiting for register read [0x{addr:04x}]",
        )
        return self.read_mmio_u8(ASMT_REG_RDATA)

    def _write_reg(self, addr: int, value: int, *, verify: bool) -> None:
        self.write_mmio_u16(ASMT_REG_ADDR, addr)
        self._poll_until(
            lambda: not (self.read_mmio_u8(ASMT_REG_STATUS) & ASMT_REG_STATUS_BUSY),
            timeout=REG_TIMEOUT_SEC,
            message=f"timed out waiting for register bus idle [0x{addr:04x}]",
        )
        self.write_mmio_u8(ASMT_REG_WDATA, value)
        self._poll_until(
            lambda: not (self.read_mmio_u8(ASMT_REG_STATUS) & ASMT_REG_STATUS_BUSY),
            timeout=REG_TIMEOUT_SEC,
            message=f"timed out waiting for register write [0x{addr:04x}] = 0x{value:02x}",
        )
        if verify:
            self._poll_until(
                lambda: self._read_reg(addr) == value,
                timeout=REG_TIMEOUT_SEC,
                message=f"timed out verifying register [0x{addr:04x}] = 0x{value:02x}",
            )

    def _pread_exact(self, offset: int, size: int) -> bytes:
        if self._config_fd is None:
            raise FlasherError("PCI config space is not open")
        try:
            data = os.pread(self._config_fd, size, offset)
        except OSError as exc:
            raise FlasherError(
                f"failed to read PCI config on {self.info.bdf} at 0x{offset:x}: {exc}"
            ) from exc
        if len(data) != size:
            raise FlasherError(
                f"short read from PCI config on {self.info.bdf}: wanted {size} bytes, got {len(data)}"
            )
        return data

    def _pwrite_exact(self, offset: int, payload: bytes) -> None:
        if self._config_fd is None:
            raise FlasherError("PCI config space is not open")
        try:
            written = os.pwrite(self._config_fd, payload, offset)
        except OSError as exc:
            raise FlasherError(
                f"failed to write PCI config on {self.info.bdf} at 0x{offset:x}: {exc}"
            ) from exc
        if written != len(payload):
            raise FlasherError(
                f"short write to PCI config on {self.info.bdf}: wrote {written}/{len(payload)} bytes"
            )

    def _bar_slice(self, offset: int, size: int) -> bytes:
        if self._bar0 is None or self._bar0_size is None:
            raise FlasherError("BAR0 is not mapped")
        if offset + size > self._bar0_size:
            raise FlasherError(
                f"BAR0 access beyond mapped range on {self.info.bdf}: 0x{offset:x}+{size}"
            )
        start = self._bar0_offset + offset
        end = start + size
        return self._bar0[start:end]

    def _bar_write(self, offset: int, payload: bytes) -> None:
        if self._bar0 is None or self._bar0_size is None:
            raise FlasherError("BAR0 is not mapped")
        if offset + len(payload) > self._bar0_size:
            raise FlasherError(
                f"BAR0 write beyond mapped range on {self.info.bdf}: 0x{offset:x}+{len(payload)}"
            )
        start = self._bar0_offset + offset
        end = start + len(payload)
        self._bar0[start:end] = payload

    def _poll_until(
        self,
        predicate,
        *,
        timeout: float,
        message: str,
        interval: float = POLL_INTERVAL_SEC,
    ) -> None:
        deadline = time.monotonic() + timeout
        while True:
            if predicate():
                return
            if time.monotonic() >= deadline:
                raise TimeoutError(message)
            time.sleep(interval)

    def _ensure_enabled(self) -> None:
        enable_path = self.info.sysfs_path / "enable"
        try:
            enable_path.write_text("1")
        except OSError:
            # Older kernels can reject writes when the device is already enabled.
            pass

    def _ensure_bar0_mapped(self) -> None:
        if self._bar0 is not None:
            return

        try:
            self._resource_fd = os.open(self.info.sysfs_path / "resource0", os.O_RDWR | os.O_SYNC)
            self._bar0_size = _bar_size(self.info.sysfs_path, index=0)
            if self._bar0_size < 0x5043:
                raise FlasherError(
                    f"BAR0 on {self.info.bdf} is too small ({self._bar0_size} bytes) for ASMedia protocol"
                )
            self._bar0, self._bar0_offset = self._open_bar0_mapping()
        except FlasherError:
            raise
        except OSError as exc:
            raise FlasherError(f"failed to open BAR0 on {self.info.bdf}: {exc}") from exc

    def _open_bar0_mapping(self) -> tuple[mmap.mmap, int]:
        if self._resource_fd is None or self._bar0_size is None:
            raise FlasherError("BAR0 resources are not open")

        try:
            return (
                mmap.mmap(
                    self._resource_fd,
                    self._bar0_size,
                    flags=mmap.MAP_SHARED,
                    prot=mmap.PROT_READ | mmap.PROT_WRITE,
                ),
                0,
            )
        except OSError as resource_exc:
            bar_start = _bar_start(self.info.sysfs_path, index=0)
            page_size = mmap.PAGESIZE
            map_base = bar_start & ~(page_size - 1)
            map_offset = bar_start - map_base
            map_length = map_offset + self._bar0_size

            try:
                self._devmem_fd = os.open("/dev/mem", os.O_RDWR | os.O_SYNC)
            except OSError as devmem_open_exc:
                raise FlasherError(
                    f"failed to map BAR0 on {self.info.bdf}: resource0 mmap failed ({resource_exc}); "
                    f"/dev/mem open failed ({devmem_open_exc})"
                ) from devmem_open_exc

            try:
                return (
                    mmap.mmap(
                        self._devmem_fd,
                        map_length,
                        flags=mmap.MAP_SHARED,
                        prot=mmap.PROT_READ | mmap.PROT_WRITE,
                        offset=map_base,
                    ),
                    map_offset,
                )
            except OSError as devmem_map_exc:
                raise FlasherError(
                    f"failed to map BAR0 on {self.info.bdf}: resource0 mmap failed ({resource_exc}); "
                    f"/dev/mem mmap failed ({devmem_map_exc})"
                ) from devmem_map_exc

    def _unbind_driver(self, driver: str) -> None:
        driver_path = Path("/sys/bus/pci/drivers") / driver / "unbind"
        try:
            driver_path.write_text(self.info.bdf)
        except OSError as exc:
            raise FlasherError(
                f"failed to unbind {self.info.bdf} from driver {driver}: {exc}"
            ) from exc

    def _bind_driver(self, driver: str) -> None:
        driver_path = Path("/sys/bus/pci/drivers") / driver / "bind"
        try:
            driver_path.write_text(self.info.bdf)
        except OSError as exc:
            raise FlasherError(
                f"failed to re-bind {self.info.bdf} to driver {driver}: {exc}"
            ) from exc


def ensure_protocol_supported(info: DeviceInfo, *, force: bool) -> None:
    if info.protocol_supported:
        return
    if force:
        return
    raise DeviceCompatibilityError(
        "the selected device is not in the known-safe protocol allowlist; "
        "pass --force-device only after you verify compatibility"
    )


def format_fw_version(value: int) -> str:
    return f"0x{value:012x}"


def _read_hex_file(path: Path) -> int:
    return int(path.read_text().strip(), 16)


def _resolve_driver_name(path: Path) -> str | None:
    driver_link = path / "driver"
    if not driver_link.exists():
        return None
    return driver_link.resolve().name


def _bar_size(device_path: Path, *, index: int) -> int:
    start, end = _bar_window(device_path, index=index)
    if end < start:
        return 0
    return end - start + 1


def _bar_start(device_path: Path, *, index: int) -> int:
    start, _end = _bar_window(device_path, index=index)
    return start


def _bar_window(device_path: Path, *, index: int) -> tuple[int, int]:
    resource_lines = (device_path / "resource").read_text().splitlines()
    start_str, end_str, _flags_str = resource_lines[index].split()
    start = int(start_str, 16)
    end = int(end_str, 16)
    return start, end
