from __future__ import annotations

import argparse
import os
from pathlib import Path
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
from .flashrom_backend import FlashromConfig, default_backup_path, read_flash, write_flash
from .internal_flash import (
    extract_permanent_firmware,
    probe_internal_spi_rom,
    read_internal_flash,
    write_internal_flash,
)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="asm3042-fw",
        description="ASMedia xHCI firmware utility for 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",
    )

    permanent_read_internal_parser = subparsers.add_parser(
        "permanent-read-internal",
        help="read the onboard SPI flash through the ASMedia internal PCIe protocol",
    )
    permanent_read_internal_parser.add_argument("output", help="path to the output backup file")
    permanent_read_internal_parser.add_argument("--bdf", help="PCI BDF like 0000:03:00.0")
    permanent_read_internal_parser.add_argument(
        "--force-device",
        action="store_true",
        help="allow using a device outside the known-safe allowlist",
    )
    permanent_read_internal_parser.add_argument(
        "--keep-driver-bound",
        action="store_true",
        help="skip the default unbind/rebind cycle around the flash access",
    )
    permanent_read_internal_parser.add_argument(
        "--size",
        type=lambda value: int(value, 0),
        help="optional byte count to read; defaults to the detected full SPI ROM size",
    )

    permanent_write_internal_parser = subparsers.add_parser(
        "permanent-write-internal",
        help="permanently write the onboard SPI flash through the ASMedia internal PCIe protocol",
    )
    permanent_write_internal_parser.add_argument("firmware", help="path to the firmware .bin file")
    permanent_write_internal_parser.add_argument("--bdf", help="PCI BDF like 0000:03:00.0")
    permanent_write_internal_parser.add_argument(
        "--force-device",
        action="store_true",
        help="allow using a device outside the known-safe allowlist",
    )
    permanent_write_internal_parser.add_argument(
        "--keep-driver-bound",
        action="store_true",
        help="skip the default unbind/rebind cycle around the flash access",
    )
    permanent_write_internal_parser.add_argument(
        "--backup-file",
        help="where to save a mandatory SPI flash backup before writing",
    )
    permanent_write_internal_parser.add_argument(
        "--skip-backup",
        action="store_true",
        help="skip the pre-write internal flash backup; not recommended",
    )
    permanent_write_internal_parser.add_argument(
        "--overwrite-backup",
        action="store_true",
        help="allow overwriting an existing backup file",
    )
    permanent_write_internal_parser.add_argument(
        "--no-verify",
        action="store_true",
        help="skip the post-write readback verification",
    )

    permanent_read_parser = subparsers.add_parser(
        "permanent-read",
        help="read the external SPI flash with flashrom and save a backup image",
    )
    permanent_read_parser.add_argument("output", help="path to the output backup file")
    permanent_read_parser.add_argument(
        "--programmer",
        required=True,
        help="flashrom programmer string, for example ch341a_spi or serprog:dev=/dev/ttyACM0",
    )
    permanent_read_parser.add_argument(
        "--chip",
        help="optional flash chip name passed to flashrom with -c",
    )
    permanent_read_parser.add_argument(
        "--flashrom-bin",
        default="flashrom",
        help="path to the flashrom executable",
    )

    permanent_write_parser = subparsers.add_parser(
        "permanent-write",
        help="write the external SPI flash permanently using flashrom and an external programmer",
    )
    permanent_write_parser.add_argument("firmware", help="path to the firmware .bin file")
    permanent_write_parser.add_argument(
        "--programmer",
        required=True,
        help="flashrom programmer string, for example ch341a_spi or serprog:dev=/dev/ttyACM0",
    )
    permanent_write_parser.add_argument(
        "--chip",
        help="optional flash chip name passed to flashrom with -c",
    )
    permanent_write_parser.add_argument(
        "--flashrom-bin",
        default="flashrom",
        help="path to the flashrom executable",
    )
    permanent_write_parser.add_argument(
        "--backup-file",
        help="where to save a mandatory SPI flash backup before writing",
    )
    permanent_write_parser.add_argument(
        "--skip-backup",
        action="store_true",
        help="skip the pre-write flash backup; not recommended",
    )
    permanent_write_parser.add_argument(
        "--overwrite-backup",
        action="store_true",
        help="allow overwriting an existing backup file",
    )
    permanent_write_parser.add_argument(
        "--no-verify",
        action="store_true",
        help="pass -n to flashrom and skip post-write verification",
    )

    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,
            )
        if args.command == "permanent-read-internal":
            _require_root(args.command)
            return _cmd_permanent_read_internal(
                args.output,
                bdf=args.bdf,
                force_device=args.force_device,
                keep_driver_bound=args.keep_driver_bound,
                size=args.size,
            )
        if args.command == "permanent-write-internal":
            _require_root(args.command)
            return _cmd_permanent_write_internal(
                args.firmware,
                bdf=args.bdf,
                force_device=args.force_device,
                keep_driver_bound=args.keep_driver_bound,
                backup_file=args.backup_file,
                skip_backup=args.skip_backup,
                overwrite_backup=args.overwrite_backup,
                verify=not args.no_verify,
            )
        if args.command == "permanent-read":
            return _cmd_permanent_read(
                args.output,
                programmer=args.programmer,
                chip=args.chip,
                flashrom_bin=args.flashrom_bin,
            )
        if args.command == "permanent-write":
            return _cmd_permanent_write(
                args.firmware,
                programmer=args.programmer,
                chip=args.chip,
                flashrom_bin=args.flashrom_bin,
                backup_file=args.backup_file,
                skip_backup=args.skip_backup,
                overwrite_backup=args.overwrite_backup,
                verify=not args.no_verify,
            )
    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 '-'}")
    try:
        permanent_image = extract_permanent_firmware(image)
    except ValueError:
        pass
    else:
        print(f"permanent-raw-size: {permanent_image.raw_size} bytes")
    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 _cmd_permanent_read(
    output_path: str,
    *,
    programmer: str,
    chip: str | None,
    flashrom_bin: str,
) -> int:
    output = Path(output_path)
    config = FlashromConfig(programmer=programmer, flashrom_bin=flashrom_bin, chip=chip)

    print(f"mode: permanent-read")
    print(f"programmer: {programmer}")
    if chip:
        print(f"chip: {chip}")
    print(f"output: {output}")

    read_flash(output, config=config)

    print("status: backup-complete")
    return 0


def _cmd_permanent_read_internal(
    output_path: str,
    *,
    bdf: str | None,
    force_device: bool,
    keep_driver_bound: bool,
    size: int | None,
) -> int:
    info = pick_device(bdf)
    ensure_protocol_supported(info, force=force_device)
    output = Path(output_path)

    print("mode: permanent-read-internal")
    print(_format_device_line(info))

    controller = AsmediaXhciController(info)
    with controller.temporarily_unbound(enabled=not keep_driver_bound):
        with controller:
            running_version = _try_get_firmware_version(controller)
            rom, data = read_internal_flash(controller, size=size)

    _print_running_version("firmware-running", running_version)
    print(_format_spi_rom_line(rom))
    print(f"output: {output}")
    print(f"bytes-read: {len(data)}")
    _write_bytes(output, data, label="internal flash backup")
    print("status: internal-backup-complete")
    return 0


def _cmd_permanent_write_internal(
    firmware_path: str,
    *,
    bdf: str | None,
    force_device: bool,
    keep_driver_bound: bool,
    backup_file: str | None,
    skip_backup: bool,
    overwrite_backup: bool,
    verify: bool,
) -> int:
    image = FirmwareImage.from_path(firmware_path)
    permanent_image = extract_permanent_firmware(image)
    firmware = Path(firmware_path)
    info = pick_device(bdf)
    ensure_protocol_supported(info, force=force_device)

    print("mode: permanent-write-internal")
    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)}")
    print(f"firmware-family: {permanent_image.family_tag}")
    print(f"firmware-raw-size: {permanent_image.raw_size} bytes")

    controller = AsmediaXhciController(info)
    with controller.temporarily_unbound(enabled=not keep_driver_bound):
        with controller:
            before = _try_get_firmware_version(controller)
            _print_running_version("firmware-running-before", before)

            if skip_backup:
                rom = probe_internal_spi_rom(controller)
                print(_format_spi_rom_line(rom))
                print("backup-file: skipped")
            else:
                backup_path = Path(backup_file) if backup_file else default_backup_path(firmware)
                if backup_path.exists() and not overwrite_backup:
                    raise FlasherError(
                        f"backup file already exists: {backup_path}; pass --overwrite-backup or choose --backup-file"
                    )
                rom, backup_data = read_internal_flash(controller)
                print(_format_spi_rom_line(rom))
                print(f"backup-file: {backup_path}")
                print(f"backup-size: {len(backup_data)} bytes")
                _write_bytes(backup_path, backup_data, label="internal flash backup")

            rom = write_internal_flash(controller, permanent_image, verify=verify)

    print(_format_spi_rom_line(rom))
    print("status: internal-permanent-write-complete")
    print("next-step: perform a cold reboot or full power cycle before validating the new SPI image")
    return 0


def _cmd_permanent_write(
    firmware_path: str,
    *,
    programmer: str,
    chip: str | None,
    flashrom_bin: str,
    backup_file: str | None,
    skip_backup: bool,
    overwrite_backup: bool,
    verify: bool,
) -> int:
    image = FirmwareImage.from_path(firmware_path)
    firmware = Path(firmware_path)
    config = FlashromConfig(programmer=programmer, flashrom_bin=flashrom_bin, chip=chip)

    print("mode: permanent-write")
    print(f"programmer: {programmer}")
    if chip:
        print(f"chip: {chip}")
    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)}")

    if not skip_backup:
        backup_path = Path(backup_file) if backup_file else default_backup_path(firmware)
        if backup_path.exists() and not overwrite_backup:
            raise FlasherError(
                f"backup file already exists: {backup_path}; pass --overwrite-backup or choose --backup-file"
            )
        print(f"backup-file: {backup_path}")
        read_flash(backup_path, config=config)
    else:
        print("backup-file: skipped")

    write_flash(firmware, config=config, verify=verify)

    print("status: permanent-write-complete")
    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 _format_spi_rom_line(rom) -> str:
    ids = f"{rom.manufacturer_id:02x}:{rom.memory_type_id:02x}:{rom.capacity_id:02x}"
    return (
        f"spi-rom: {rom.manufacturer} {rom.model} "
        f"(id={ids}, size={rom.flash_size_bytes} bytes, sector={rom.sector_size_bytes} bytes)"
    )


def _write_bytes(path: Path, data: bytes, *, label: str) -> None:
    try:
        path.write_bytes(data)
    except OSError as exc:
        raise FlasherError(f"failed to write {label} to {path}: {exc}") from exc


def _try_get_firmware_version(controller: AsmediaXhciController) -> int | None:
    try:
        return controller.get_firmware_version()
    except (FlasherError, TimeoutError):
        return None


def _print_running_version(label: str, value: int | None) -> None:
    if value is None:
        print(f"{label}: unavailable")
        return
    print(f"{label}: {format_fw_version(value)}")


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


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