from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import re


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


@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]

    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)
