from __future__ import annotations

import binascii
import configparser
from dataclasses import dataclass
from pathlib import Path
import re
import struct


MAX_FIRMWARE_SIZE = 128 * 1024
HEADER_TAG_RE = re.compile(rb"([A-Z0-9]{4,8}_(?:RCFG|FW))")

FAMILY_DEVICE_TYPES: dict[str, int] = {
    "U2104": 0,
    "2104B": 1,
    "2114A": 1,
    "2214A": 2,
    "1343A": 3,
    "3306A": 4,
    "3306B": 4,
    "2324A": 5,
}

CFG_RECORD_ADDRESSES: dict[int, dict[str, int]] = {
    0: {"ssid": 0xF392, "svid": 0xF390},
    1: {"ssid": 0xE262, "svid": 0xE260},
    2: {"ssid": 0x52, "svid": 0x50, "pcie_speed": 0x10063},
    4: {"ssid": 0x52, "svid": 0x50},
    5: {"ssid": 0x5A, "svid": 0x58, "pcie_speed": 0xAE},
}

FW_SIGNATURE_SIZE = 8
FW_CHECKSUM_SIZE = 1
FW_CRC_SIZE = 4
FW_EXTRA_SIZE_BY_TYPE = {
    5: 0x41,
}


@dataclass(frozen=True)
class PermanentConfigRecord:
    width: int
    address: int
    value: int

    def to_bytes(self) -> bytes:
        if self.width not in {0, 1, 2, 4, 8}:
            raise ValueError(f"unsupported config record width: {self.width}")
        if not 0 <= self.address <= 0xFFFF:
            raise ValueError(f"config record address is out of range: 0x{self.address:x}")
        if not 0 <= self.value <= 0xFFFFFFFF:
            raise ValueError(f"config record value is out of range: 0x{self.value:x}")

        payload = bytearray(8)
        payload[0] = 0xCC
        payload[1] = self.width
        payload[2:4] = struct.pack("<H", self.address)
        if self.width >= 1:
            payload[4] = self.value & 0xFF
        if self.width >= 2:
            payload[5] = (self.value >> 8) & 0xFF
        if self.width >= 4:
            payload[6] = (self.value >> 16) & 0xFF
        if self.width >= 8:
            payload[7] = (self.value >> 24) & 0xFF
        return bytes(payload)


@dataclass(frozen=True)
class MpToolIniConfig:
    source: str
    subsystem_vendor_id: int | None = None
    subsystem_device_id: int | None = None
    pcie_link_speed: int | None = None

    @property
    def has_overrides(self) -> bool:
        return any(
            value is not None
            for value in (self.subsystem_vendor_id, self.subsystem_device_id, self.pcie_link_speed)
        )

    def build_records(self, *, device_type: int) -> tuple[PermanentConfigRecord, ...]:
        address_map = CFG_RECORD_ADDRESSES.get(device_type, {})
        records: list[PermanentConfigRecord] = []

        if self.subsystem_device_id is not None:
            address = address_map.get("ssid")
            if address is None:
                raise ValueError(
                    f"device type {device_type} does not expose a Windows-equivalent SSID config register"
                )
            records.append(
                PermanentConfigRecord(width=2, address=address, value=self.subsystem_device_id)
            )

        if self.subsystem_vendor_id is not None:
            address = address_map.get("svid")
            if address is None:
                raise ValueError(
                    f"device type {device_type} does not expose a Windows-equivalent SVID config register"
                )
            records.append(
                PermanentConfigRecord(width=2, address=address, value=self.subsystem_vendor_id)
            )

        if self.pcie_link_speed is not None:
            address = address_map.get("pcie_speed")
            if address is None:
                raise ValueError(
                    f"device type {device_type} does not expose a Windows-equivalent PCIe speed config register"
                )
            records.append(
                PermanentConfigRecord(width=1, address=address, value=self.pcie_link_speed)
            )

        return tuple(records)


@dataclass(frozen=True)
class PermanentFirmwareImage:
    source: str
    family_tag: str
    device_type: int
    config_length: int
    raw_data: bytes
    rom_data: bytes
    config_records: tuple[PermanentConfigRecord, ...]

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

    @property
    def rom_size(self) -> int:
        return len(self.rom_data)


@dataclass(frozen=True)
class FirmwareImage:
    source: str
    data: bytes
    header_tags: tuple[str, ...]

    @classmethod
    def from_path(cls, path: str | Path) -> "FirmwareImage":
        firmware_path = Path(path)
        return cls.from_bytes(firmware_path.read_bytes(), source=str(firmware_path))

    @classmethod
    def from_bytes(cls, data: bytes, source: str = "<memory>") -> "FirmwareImage":
        if not data:
            raise ValueError("firmware image is empty")
        if len(data) % 2:
            raise ValueError("firmware image size must be even")
        if len(data) > MAX_FIRMWARE_SIZE:
            raise ValueError(
                f"firmware image is too large ({len(data)} bytes > {MAX_FIRMWARE_SIZE} bytes)"
            )

        tags: list[str] = []
        seen: set[str] = set()
        for match in HEADER_TAG_RE.finditer(data[:512]):
            tag = match.group(1).decode("ascii")
            if tag not in seen:
                tags.append(tag)
                seen.add(tag)

        return cls(source=source, data=data, header_tags=tuple(tags))

    @property
    def size(self) -> int:
        return len(self.data)

    @property
    def word_count(self) -> int:
        return len(self.data) // 2

    @property
    def family_tag(self) -> str | None:
        if not self.header_tags:
            return None
        return self.header_tags[0].split("_", maxsplit=1)[0]

    @property
    def config_length(self) -> int | None:
        if len(self.data) < 6:
            return None
        return int.from_bytes(self.data[4:6], "little")

    def build_permanent_image(
        self,
        *,
        config_records: tuple[PermanentConfigRecord, ...] = (),
    ) -> PermanentFirmwareImage:
        family_tag = self.family_tag
        if not family_tag or family_tag not in FAMILY_DEVICE_TYPES:
            raise ValueError(
                "internal permanent flashing currently supports only ASMedia images "
                "with known RCFG/FW headers"
            )

        config_length = self.config_length
        if config_length is None:
            raise ValueError("firmware image is truncated before the config size field")
        if config_length < 0x10:
            raise ValueError(f"firmware config length is too small: 0x{config_length:x}")
        if config_length + 5 > len(self.data):
            raise ValueError("firmware image is truncated before the config checksum/CRC fields")

        expected_cfg_checksum = sum(self.data[:config_length]) & 0xFF
        actual_cfg_checksum = self.data[config_length]
        if actual_cfg_checksum != expected_cfg_checksum:
            raise ValueError(
                f"firmware config checksum mismatch: expected 0x{expected_cfg_checksum:02x}, "
                f"got 0x{actual_cfg_checksum:02x}"
            )

        device_type = FAMILY_DEVICE_TYPES[family_tag]
        size_width = 2 if device_type in {0, 1} else 4
        raw_size_offset = config_length + 5
        raw_offset = raw_size_offset + size_width
        raw_size_end = raw_size_offset + size_width
        if raw_size_end > len(self.data):
            raise ValueError("firmware image is truncated before the raw firmware size field")

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

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

        tail_offset = raw_end
        tail_size = FW_SIGNATURE_SIZE + FW_CHECKSUM_SIZE + FW_CRC_SIZE
        extra_size = FW_EXTRA_SIZE_BY_TYPE.get(device_type, 0)
        tail_end = tail_offset + tail_size + extra_size
        if tail_end > len(self.data):
            raise ValueError("firmware image is truncated before the permanent ROM trailer")

        fw_signature = self.data[tail_offset : tail_offset + FW_SIGNATURE_SIZE]
        fw_checksum = self.data[tail_offset + FW_SIGNATURE_SIZE]
        fw_crc = self.data[
            tail_offset + FW_SIGNATURE_SIZE + FW_CHECKSUM_SIZE : tail_offset + tail_size
        ]
        fw_extra = self.data[tail_offset + tail_size : tail_end]

        customer_info = self.data[0x10:0x30] if device_type > 0 else b""
        config_blob = bytearray(self.data[:0x10])
        config_blob.extend(self.data[: config_length - 0x10])
        config_blob.extend(customer_info)
        for record in config_records:
            config_blob.extend(record.to_bytes())

        new_config_length = config_length + len(customer_info) + len(config_records) * 8
        config_blob[4:6] = struct.pack("<H", new_config_length)

        config_checksum = sum(config_blob[:new_config_length]) & 0xFF
        config_blob.append(config_checksum)
        config_crc = binascii.crc32(config_blob[: new_config_length + 1]) & 0xFFFFFFFF
        config_blob.extend(struct.pack("<I", config_crc))

        rom_data = bytearray(config_blob)
        rom_data.extend(raw_size.to_bytes(size_width, "little"))
        rom_data.extend(self.data[raw_offset:raw_end])
        rom_data.extend(fw_signature)
        rom_data.append(fw_checksum)
        rom_data.extend(fw_crc)
        rom_data.extend(fw_extra)

        return PermanentFirmwareImage(
            source=self.source,
            family_tag=family_tag,
            device_type=device_type,
            config_length=new_config_length,
            raw_data=self.data[raw_offset:raw_end],
            rom_data=bytes(rom_data),
            config_records=config_records,
        )

    def iter_sram_writes(self):
        """Yield (SRAM address, 32-bit word) pairs in ASMedia interleaved order."""
        words = self.word_count
        index = 0
        addr = 0

        while index < words:
            data = self._word_at(index)
            upper_index = index | 0x4000
            if upper_index < words:
                data |= self._word_at(upper_index) << 16

            yield addr, data

            index += 1
            if index & 0x4000:
                index += 0x4000
            addr += 2

    def _word_at(self, index: int) -> int:
        offset = index * 2
        return self.data[offset] | (self.data[offset + 1] << 8)


def parse_mptool_ini(path: str | Path) -> MpToolIniConfig:
    ini_path = Path(path)
    parser = configparser.ConfigParser()
    parser.optionxform = str.lower
    try:
        with ini_path.open("r", encoding="utf-8-sig") as handle:
            parser.read_file(handle)
    except OSError as exc:
        raise ValueError(f"failed to read MPTool ini file {ini_path}: {exc}") from exc

    config_section = parser["config"] if parser.has_section("config") else {}
    subsystem_enabled = _ini_bool(config_section.get("subsys_enable", "0"))
    subsystem_vendor_id: int | None = None
    subsystem_device_id: int | None = None
    pcie_link_speed: int | None = None

    if subsystem_enabled:
        if "svid" not in config_section or "ssid" not in config_section:
            raise ValueError("MPTool ini enables subsystem override but is missing svid/ssid")
        subsystem_vendor_id = _parse_ini_int(config_section["svid"], label="svid")
        subsystem_device_id = _parse_ini_int(config_section["ssid"], label="ssid")

    if "pcispeed" in config_section:
        pcie_link_speed = _parse_ini_int(config_section["pcispeed"], label="pcispeed")

    return MpToolIniConfig(
        source=str(ini_path),
        subsystem_vendor_id=subsystem_vendor_id,
        subsystem_device_id=subsystem_device_id,
        pcie_link_speed=pcie_link_speed,
    )


def _ini_bool(value: str) -> bool:
    return value.strip().lower() in {"1", "true", "yes", "on"}


def _parse_ini_int(value: str, *, label: str) -> int:
    text = value.strip()
    if not text:
        raise ValueError(f"empty integer value for {label}")
    try:
        return int(text, 16 if text.startswith(("0x", "0X")) else 16 if _looks_like_hex(text) else 10)
    except ValueError as exc:
        raise ValueError(f"invalid integer value for {label}: {value!r}") from exc


def _looks_like_hex(value: str) -> bool:
    return all(char in "0123456789abcdefABCDEF" for char in value)
