from __future__ import annotations

from dataclasses import dataclass
import time

from .device import AsmediaXhciController, FlasherError
from .firmware import FirmwareImage
from .spi_rom_table import SPI_ROM_PREFIXES


FLASH_CHUNK_SIZE = 0x10000
READ_CHUNK_SIZE = 8
READ_READY_MASK = 1 << 0
WRITE_READY_MASK = 1 << 1
CONTROL_OFFSET = 0xE0
DATA_READ0_OFFSET = 0xF0
DATA_READ1_OFFSET = 0xF4
DATA_WRITE0_OFFSET = 0xF8
DATA_WRITE1_OFFSET = 0xFC

DEFAULT_POLL_INTERVAL_SEC = 0.0005
SHORT_TIMEOUT_SEC = 0.25
MAX_READY_TIMEOUT_SEC = 20.0

INTERNAL_DEVICE_TYPES = {
    0x3042: 2,
    0x3142: 2,
}

FAMILY_LAYOUTS: dict[str, tuple[int, int]] = {
    "2104B": (7, 2),
    "2114A": (7, 2),
    "2214A": (9, 4),
    "2324A": (9, 4),
    "3306A": (9, 4),
    "3306B": (9, 4),
    "1343A": (9, 4),
    "U2104": (7, 2),
}


@dataclass(frozen=True)
class SpiRomInfo:
    prefix: bytes
    manufacturer: str
    model: str

    @property
    def read_id_command(self) -> int:
        return self.prefix[0]

    @property
    def manufacturer_id(self) -> int:
        return self.prefix[8]

    @property
    def memory_type_id(self) -> int:
        return self.prefix[9]

    @property
    def capacity_id(self) -> int:
        return self.prefix[10]

    @property
    def rom_size_units(self) -> int:
        return int.from_bytes(self.prefix[4:6], "little")

    @property
    def flash_size_bytes(self) -> int:
        return self.rom_size_units << 12

    @property
    def sector_size_bytes(self) -> int:
        return self.prefix[3] << 12

    @property
    def erase_descriptor_low(self) -> int:
        return int.from_bytes(self.prefix[0:4], "little")

    @property
    def erase_descriptor_high(self) -> int:
        return (
            int.from_bytes(self.prefix[4:6], "little")
            + (self.prefix[6] << 8)
            + (self.prefix[7] << 16)
            + (self.prefix[8] << 24)
        )

    @property
    def ids(self) -> tuple[int, int, int]:
        return (self.manufacturer_id, self.memory_type_id, self.capacity_id)

    @property
    def requires_sst_write_mode(self) -> bool:
        return self.manufacturer_id == 0xBF

    @property
    def uses_split_erase(self) -> bool:
        return (
            (self.manufacturer_id == 0xC8 and self.memory_type_id == 0x40)
            or (self.manufacturer_id == 0x20 and self.memory_type_id == 0x20)
            or (self.manufacturer_id == 0x68 and self.memory_type_id == 0x40)
        )


@dataclass(frozen=True)
class PermanentFirmwareImage:
    source: str
    family_tag: str
    raw_data: bytes

    @property
    def raw_size(self) -> int:
        return len(self.raw_data)


SPI_ROM_INFOS: tuple[SpiRomInfo, ...] = tuple(
    SpiRomInfo(prefix=prefix, manufacturer=manufacturer, model=model)
    for prefix, manufacturer, model in SPI_ROM_PREFIXES
)


def probe_internal_spi_rom(controller: AsmediaXhciController) -> SpiRomInfo:
    device_type = INTERNAL_DEVICE_TYPES.get(controller.info.device_id)
    if device_type is None:
        raise FlasherError(
            f"internal permanent flashing is not supported for device 0x{controller.info.device_id:04x}"
        )

    _set_spi_page_size(controller, 0x80)
    rom = _find_spi_rom_info(controller)

    if rom.manufacturer_id == 0xC2 and rom.memory_type_id == 0x22 and rom.capacity_id in {0x10, 0x11}:
        _set_spi_page_size(controller, 0x20)

    if device_type in {2, 4} and rom.flash_size_bytes < 0x20000:
        raise FlasherError(
            f"detected SPI ROM is too small ({rom.flash_size_bytes} bytes) for ASM3042/ASM3142"
        )

    return rom


def extract_permanent_firmware(image: FirmwareImage) -> PermanentFirmwareImage:
    family_tag = image.family_tag
    if not family_tag or family_tag not in FAMILY_LAYOUTS:
        raise ValueError(
            "internal permanent flashing currently supports only ASMedia images "
            "with known RCFG/FW headers"
        )

    header_len = image.config_length
    if header_len is None:
        raise ValueError("firmware image is truncated before the config size field")
    raw_offset, size_width = FAMILY_LAYOUTS[family_tag]
    size_start = header_len + raw_offset - size_width
    size_end = size_start + size_width
    data_start = header_len + raw_offset

    if size_end > len(image.data):
        raise ValueError("firmware image header is truncated before the raw firmware size field")

    raw_size = int.from_bytes(image.data[size_start:size_end], "little")
    if raw_size <= 0:
        raise ValueError("firmware image raw size is zero")

    data_end = data_start + raw_size
    if data_end > len(image.data):
        raise ValueError(
            f"firmware image raw payload overruns the file ({data_end} > {len(image.data)})"
        )

    return PermanentFirmwareImage(
        source=image.source,
        family_tag=family_tag,
        raw_data=image.data[data_start:data_end],
    )


def write_internal_flash(
    controller: AsmediaXhciController,
    firmware: PermanentFirmwareImage,
    *,
    verify: bool = True,
) -> SpiRomInfo:
    rom = _init_spi_rom(controller)
    if rom.sector_size_bytes <= 0:
        raise FlasherError("detected SPI ROM reports an invalid erase sector size")
    erase_size = _round_up(firmware.raw_size, rom.sector_size_bytes)

    if erase_size > rom.flash_size_bytes:
        raise FlasherError(
            f"firmware payload ({erase_size} bytes after sector alignment) does not fit into detected SPI ROM "
            f"({rom.flash_size_bytes} bytes)"
        )

    for offset in range(0, erase_size, FLASH_CHUNK_SIZE):
        chunk_erase_size = min(FLASH_CHUNK_SIZE, erase_size - offset)
        _erase_region(controller, rom, offset=offset, size=chunk_erase_size)
        last_error: Exception | None = None
        chunk = firmware.raw_data[offset : offset + FLASH_CHUNK_SIZE]
        if not chunk:
            continue
        for retry in range(3):
            try:
                _write_section(
                    controller,
                    rom,
                    chunk,
                    offset=offset,
                    retry_index=retry,
                )
                last_error = None
                break
            except Exception as exc:  # pragma: no cover - hardware dependent
                last_error = exc
                time.sleep(0.2)
        if last_error is not None:
            raise FlasherError(
                f"failed to write SPI chunk at 0x{offset:05x}: {last_error}"
            ) from last_error

    if verify:
        read_back = _read_flash(controller, firmware.raw_size)
        if read_back != firmware.raw_data:
            mismatch = next(
                idx
                for idx, (left, right) in enumerate(zip(read_back, firmware.raw_data))
                if left != right
            )
            raise FlasherError(
                f"SPI verification failed at offset 0x{mismatch:05x}: "
                f"read 0x{read_back[mismatch]:02x}, expected 0x{firmware.raw_data[mismatch]:02x}"
            )

    return rom


def read_internal_flash(
    controller: AsmediaXhciController,
    *,
    size: int | None = None,
) -> tuple[SpiRomInfo, bytes]:
    rom = probe_internal_spi_rom(controller)
    read_size = rom.flash_size_bytes if size is None else size
    if read_size <= 0:
        raise FlasherError(f"read size must be positive, got {read_size}")
    if read_size > rom.flash_size_bytes:
        raise FlasherError(
            f"requested read size ({read_size} bytes) exceeds detected SPI ROM size "
            f"({rom.flash_size_bytes} bytes)"
        )
    return rom, _read_flash(controller, read_size)


def _init_spi_rom(controller: AsmediaXhciController) -> SpiRomInfo:
    rom = probe_internal_spi_rom(controller)
    _set_write_protect(controller, enable=False)
    return rom


def _find_spi_rom_info(controller: AsmediaXhciController) -> SpiRomInfo:
    read_cache: dict[int, tuple[int, int, int]] = {}

    for rom in SPI_ROM_INFOS:
        command = rom.read_id_command
        if command not in read_cache:
            read_cache[command] = _read_spi_rom_id(controller, command)

        observed = read_cache[command]
        expected = rom.ids
        if observed[0] != expected[0] or observed[1] != expected[1]:
            continue
        if expected[2] != 0xFF and observed[2] != expected[2]:
            continue
        return rom

    attempted = ", ".join(
        f"0x{command:02x}={ids[0]:02x}:{ids[1]:02x}:{ids[2]:02x}"
        for command, ids in sorted(read_cache.items())
    )
    raise FlasherError(f"unsupported SPI ROM ID ({attempted or 'no-response'})")


def _read_spi_rom_id(
    controller: AsmediaXhciController,
    command: int,
) -> tuple[int, int, int]:
    data = _read_memory(controller, mem_type=0x20, size=3, addr=0, sec_low=command)
    return (data[0], data[1], data[2])


def _set_spi_page_size(controller: AsmediaXhciController, page_size: int) -> None:
    _write_command(
        controller,
        cmd=0x1B,
        mem_type=0x10,
        size=1,
        addr=0,
        sec_low=page_size,
        sec_high=0,
    )


def _set_write_protect(controller: AsmediaXhciController, *, enable: bool) -> tuple[int, int]:
    protect_bits = 0x9C if enable else 0x00

    _write_command(
        controller,
        cmd=0x19,
        mem_type=0x10,
        size=1,
        addr=0,
        sec_low=0x06,
        sec_high=0,
    )
    _write_status_register(controller, 0x301 | (protect_bits << 8))

    if enable:
        _write_command(
            controller,
            cmd=0x19,
            mem_type=0x10,
            size=1,
            addr=0,
            sec_low=0x04,
            sec_high=0,
        )

    return _read_status_registers(controller)


def _write_status_register(controller: AsmediaXhciController, status_word: int) -> None:
    _write_command(
        controller,
        cmd=0x1C,
        mem_type=0x10,
        size=1,
        addr=0,
        sec_low=status_word,
        sec_high=0,
    )


def _read_status_registers(controller: AsmediaXhciController) -> tuple[int, int]:
    _set_write_cmd_all(
        controller,
        cmd=0x4C,
        mem_type=0x10,
        size=2,
        addr=0,
        sec_low=0x05,
        sec_high=0,
    )
    _wait_read_ready(controller)
    data = _read_response_bytes(controller, size=2)
    return (data[0], data[1])


def _erase_region(
    controller: AsmediaXhciController,
    rom: SpiRomInfo,
    *,
    offset: int,
    size: int,
) -> None:
    if size <= 0:
        return
    if rom.sector_size_bytes == 0:
        raise FlasherError("detected SPI ROM reports zero erase sector size")
    if size % rom.sector_size_bytes:
        raise FlasherError(
            f"erase size 0x{size:x} is not aligned to the SPI sector size 0x{rom.sector_size_bytes:x}"
        )

    erase_low = rom.erase_descriptor_low
    erase_high = rom.erase_descriptor_high

    if rom.uses_split_erase:
        total_sectors = size // rom.sector_size_bytes
        sector_count = total_sectors // 2
        if sector_count <= 0:
            raise FlasherError("split erase computed an empty sector count")
        block_size = sector_count * rom.sector_size_bytes
        current = offset
        while current < offset + size:
            _write_command(
                controller,
                cmd=0x30,
                mem_type=0x10,
                size=sector_count,
                addr=current,
                sec_low=erase_low,
                sec_high=erase_high,
                wait_units=0x3E8,
            )
            time.sleep(0.05)
            current += block_size
        return

    _write_command(
        controller,
        cmd=0x30,
        mem_type=0x10,
        size=size // rom.sector_size_bytes,
        addr=offset,
        sec_low=erase_low,
        sec_high=erase_high,
        wait_units=0x7D0,
    )


def _write_section(
    controller: AsmediaXhciController,
    rom: SpiRomInfo,
    chunk: bytes,
    *,
    offset: int,
    retry_index: int,
) -> None:
    if not chunk or len(chunk) > FLASH_CHUNK_SIZE:
        raise FlasherError(
            f"internal write chunks must be between 1 and 0x{FLASH_CHUNK_SIZE:x} bytes, got 0x{len(chunk):x}"
        )

    if rom.requires_sst_write_mode:
        control = _read_memory(controller, mem_type=0x04, size=1, addr=0xF341, sec_low=0)[0]
        mode_base = retry_index << 7
        start_value = mode_base + (0x80 if (control & 0x02) else 0xF0)
        _write_command(
            controller,
            cmd=0x1D,
            mem_type=0x10,
            size=0,
            addr=offset,
            sec_low=start_value,
            sec_high=0,
        )
    else:
        _write_command(
            controller,
            cmd=0x11,
            mem_type=0x10,
            size=0,
            addr=offset,
            sec_low=0,
            sec_high=0,
        )

    for block_offset in range(0, len(chunk), READ_CHUNK_SIZE):
        block = chunk[block_offset : block_offset + READ_CHUNK_SIZE]
        low = int.from_bytes(block[:4], "little")
        high = int.from_bytes(block[4:], "little")
        command = 0x1E if len(block) == READ_CHUNK_SIZE else 0x1F
        _write_command(
            controller,
            cmd=command,
            mem_type=0x10,
            size=len(block),
            addr=offset + block_offset,
            sec_low=low,
            sec_high=high,
        )


def _read_flash(controller: AsmediaXhciController, size: int) -> bytes:
    data = bytearray()
    offset = 0
    while offset < size:
        chunk_size = min(READ_CHUNK_SIZE, size - offset)
        data.extend(_read_memory(controller, mem_type=0x10, size=chunk_size, addr=offset, sec_low=0))
        offset += chunk_size
    return bytes(data)


def _read_memory(
    controller: AsmediaXhciController,
    *,
    mem_type: int,
    size: int,
    addr: int,
    sec_low: int,
) -> bytes:
    if not 0 < size <= READ_CHUNK_SIZE:
        raise FlasherError(f"read size must be between 1 and {READ_CHUNK_SIZE}, got {size}")

    _set_write_cmd_all(
        controller,
        cmd=0x40,
        mem_type=mem_type,
        size=size,
        addr=addr,
        sec_low=sec_low,
        sec_high=0,
    )
    _wait_read_ready(controller)
    return _read_response_bytes(controller, size=size)


def _write_command(
    controller: AsmediaXhciController,
    *,
    cmd: int,
    mem_type: int,
    size: int,
    addr: int,
    sec_low: int,
    sec_high: int,
    wait_units: int = 0x14,
) -> None:
    _set_write_cmd_all(
        controller,
        cmd=cmd,
        mem_type=mem_type,
        size=size,
        addr=addr,
        sec_low=sec_low,
        sec_high=sec_high,
        wait_units=wait_units,
    )


def _set_write_cmd_all(
    controller: AsmediaXhciController,
    *,
    cmd: int,
    mem_type: int,
    size: int,
    addr: int,
    sec_low: int,
    sec_high: int,
    wait_units: int = 0x14,
) -> None:
    value_low = ((size & 0xFFFF) << 16) | ((mem_type & 0xFF) << 8) | (cmd & 0xFF)
    _write_request(controller, value_low=value_low, value_high=addr, wait_units=wait_units)
    _write_request(controller, value_low=sec_low, value_high=sec_high, wait_units=wait_units)


def _write_request(
    controller: AsmediaXhciController,
    *,
    value_low: int,
    value_high: int,
    wait_units: int,
) -> None:
    controller.write_config_u32(DATA_WRITE0_OFFSET, value_low & 0xFFFFFFFF)
    controller.write_config_u32(DATA_WRITE1_OFFSET, value_high & 0xFFFFFFFF)
    controller.write_config_u8(CONTROL_OFFSET, WRITE_READY_MASK)
    _wait_write_ready(controller, wait_units=wait_units)


def _read_response_bytes(controller: AsmediaXhciController, *, size: int) -> bytes:
    low = controller.read_config_u32(DATA_READ0_OFFSET)
    high = controller.read_config_u32(DATA_READ1_OFFSET)
    controller.write_config_u8(CONTROL_OFFSET, READ_READY_MASK)
    return (low.to_bytes(4, "little") + high.to_bytes(4, "little"))[:size]


def _wait_read_ready(controller: AsmediaXhciController) -> None:
    _poll_control_clear(
        controller,
        mask=READ_READY_MASK,
        timeout=SHORT_TIMEOUT_SEC,
        label="read-ready",
    )


def _wait_write_ready(controller: AsmediaXhciController, *, wait_units: int) -> None:
    timeout = min(MAX_READY_TIMEOUT_SEC, max(SHORT_TIMEOUT_SEC, wait_units / 100.0))
    _poll_control_clear(
        controller,
        mask=WRITE_READY_MASK,
        timeout=timeout,
        label="write-ready",
    )


def _poll_control_clear(
    controller: AsmediaXhciController,
    *,
    mask: int,
    timeout: float,
    label: str,
) -> None:
    deadline = time.monotonic() + timeout
    while True:
        value = controller.read_config_u8(CONTROL_OFFSET)
        if value != 0xFF and not (value & mask):
            return
        if time.monotonic() >= deadline:
            raise FlasherError(
                f"timed out waiting for SPI controller {label} state on {controller.info.bdf}"
            )
        time.sleep(DEFAULT_POLL_INTERVAL_SEC)


def _round_up(value: int, multiple: int) -> int:
    if value % multiple == 0:
        return value
    return value + (multiple - (value % multiple))
