from __future__ import annotations

import argparse
import os
import sys

from .device import (
    AsmediaXhciController,
    DeviceCompatibilityError,
    DeviceInfo,
    DeviceNotFoundError,
    FlasherError,
    ROM_FW_VERSION,
    discover_devices,
    ensure_protocol_supported,
    format_fw_version,
    pick_device,
)
from .firmware import FirmwareImage


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="asm3042-fw",
        description=(
            "Upload ASMedia xHCI firmware to an ASM3042-class PCIe controller under Linux"
        ),
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    subparsers.add_parser("discover", help="list ASMedia xHCI devices found in sysfs")

    inspect_parser = subparsers.add_parser("inspect", help="inspect a firmware image")
    inspect_parser.add_argument("firmware", help="path to the firmware .bin file")

    version_parser = subparsers.add_parser(
        "version", help="read the current firmware version from the controller"
    )
    version_parser.add_argument("--bdf", help="PCI BDF like 0000:03:00.0")
    version_parser.add_argument(
        "--force-device",
        action="store_true",
        help="allow using a device outside the known-safe allowlist",
    )

    upload_parser = subparsers.add_parser(
        "upload", help="upload a firmware image into controller SRAM"
    )
    upload_parser.add_argument("firmware", help="path to the firmware .bin file")
    upload_parser.add_argument("--bdf", help="PCI BDF like 0000:03:00.0")
    upload_parser.add_argument(
        "--force-device",
        action="store_true",
        help="allow using a device outside the known-safe allowlist",
    )
    upload_parser.add_argument(
        "--force-running-firmware",
        action="store_true",
        help="allow upload even if the controller already reports non-ROM firmware",
    )
    upload_parser.add_argument(
        "--keep-driver-bound",
        action="store_true",
        help="skip the default unbind/rebind cycle around the upload",
    )

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    try:
        if args.command == "discover":
            return _cmd_discover()
        if args.command == "inspect":
            return _cmd_inspect(args.firmware)
        if args.command == "version":
            _require_root(args.command)
            return _cmd_version(args.bdf, force_device=args.force_device)
        if args.command == "upload":
            _require_root(args.command)
            return _cmd_upload(
                args.firmware,
                bdf=args.bdf,
                force_device=args.force_device,
                force_running_firmware=args.force_running_firmware,
                keep_driver_bound=args.keep_driver_bound,
            )
    except (DeviceNotFoundError, DeviceCompatibilityError, FlasherError, TimeoutError, ValueError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1

    parser.error(f"unsupported command: {args.command}")
    return 2


def _cmd_discover() -> int:
    devices = discover_devices()
    if not devices:
        print("No ASMedia xHCI devices found.")
        return 1

    for device in devices:
        known = "yes" if device.protocol_supported else "no"
        driver = device.driver or "-"
        print(
            f"{device.bdf}  vendor=0x{device.vendor_id:04x}  device=0x{device.device_id:04x}  "
            f"class=0x{device.class_code:06x}  driver={driver}  known-protocol={known}"
        )
    return 0


def _cmd_inspect(firmware_path: str) -> int:
    image = FirmwareImage.from_path(firmware_path)
    print(f"source: {image.source}")
    print(f"size: {image.size} bytes ({image.word_count} words)")
    print(f"family-tag: {image.family_tag or '-'}")
    print(f"header-tags: {', '.join(image.header_tags) if image.header_tags else '-'}")
    return 0


def _cmd_version(bdf: str | None, *, force_device: bool) -> int:
    info = pick_device(bdf)
    ensure_protocol_supported(info, force=force_device)

    with AsmediaXhciController(info) as controller:
        version = controller.get_firmware_version()

    print(_format_device_line(info))
    print(f"firmware-version: {format_fw_version(version)}")
    print(f"is-rom-firmware: {'yes' if version == ROM_FW_VERSION else 'no'}")
    return 0


def _cmd_upload(
    firmware_path: str,
    *,
    bdf: str | None,
    force_device: bool,
    force_running_firmware: bool,
    keep_driver_bound: bool,
) -> int:
    image = FirmwareImage.from_path(firmware_path)
    info = pick_device(bdf)
    ensure_protocol_supported(info, force=force_device)

    print(_format_device_line(info))
    print(f"firmware-image: {image.source}")
    print(f"firmware-size: {image.size} bytes")
    if image.header_tags:
        print(f"firmware-tags: {', '.join(image.header_tags)}")

    controller = AsmediaXhciController(info)
    with controller.temporarily_unbound(enabled=not keep_driver_bound):
        with controller:
            before = controller.get_firmware_version()
            print(f"firmware-before: {format_fw_version(before)}")
            if before != ROM_FW_VERSION and not force_running_firmware:
                raise FlasherError(
                    "controller already reports non-ROM firmware; pass "
                    "--force-running-firmware to overwrite the running image"
                )

            controller.upload_firmware(image)
            after = controller.get_firmware_version()

    print(f"firmware-after: {format_fw_version(after)}")
    print("status: uploaded-to-sram")
    return 0


def _format_device_line(info: DeviceInfo) -> str:
    driver = info.driver or "-"
    return (
        f"device: {info.bdf} "
        f"(vendor=0x{info.vendor_id:04x}, device=0x{info.device_id:04x}, driver={driver})"
    )


def _require_root(command: str) -> None:
    if os.geteuid() != 0:
        raise FlasherError(f"'{command}' requires root privileges")


if __name__ == "__main__":
    raise SystemExit(main())
