from __future__ import annotations

import struct
from types import SimpleNamespace

import pytest

from asm3042_flasher.firmware import FirmwareImage
from asm3042_flasher.internal_flash import SPI_ROM_INFOS, _read_memory, extract_permanent_firmware


def test_extracts_2214a_raw_payload() -> None:
    raw_payload = bytes(range(32))
    image = FirmwareImage.from_bytes(_build_2214a_image(raw_payload), source="sample.bin")

    permanent_image = extract_permanent_firmware(image)

    assert permanent_image.source == "sample.bin"
    assert permanent_image.family_tag == "2214A"
    assert permanent_image.raw_size == len(raw_payload)
    assert permanent_image.raw_data == raw_payload


def test_extract_rejects_unknown_family_layout() -> None:
    image = FirmwareImage.from_bytes(
        b"\x00\x00\x00\x00ZZZZ_RCFG\x00ZZZZ_FW\x00\x00\x00",
        source="unknown.bin",
    )

    with pytest.raises(ValueError, match="known RCFG/FW headers"):
        extract_permanent_firmware(image)


def test_spi_rom_table_contains_known_winbond_part() -> None:
    rom = next(item for item in SPI_ROM_INFOS if item.ids == (0xEF, 0x30, 0x11))

    assert rom.manufacturer == "WinBond"
    assert rom.flash_size_bytes == 0x20000
    assert rom.sector_size_bytes == 0x1000


def test_spi_rom_table_marks_sst_write_mode() -> None:
    rom = next(item for item in SPI_ROM_INFOS if item.ids == (0xBF, 0x48, 0xBF))

    assert rom.requires_sst_write_mode is True


def test_read_memory_uses_second_read_ready_phase() -> None:
    class FakeController:
        def __init__(self) -> None:
            self.info = SimpleNamespace(bdf="0000:01:00.0")
            self._control_reads = iter([0, 0, 0x01, 0x01])
            self._data_reads = iter([
                0x00000008,
                0x00000000,
                0x001130EF,
                0x00000000,
            ])
            self.control_writes: list[int] = []

        def read_config_u8(self, offset: int) -> int:
            assert offset == 0xE0
            return next(self._control_reads)

        def read_config_u32(self, offset: int) -> int:
            assert offset in {0xF0, 0xF4}
            return next(self._data_reads)

        def write_config_u32(self, offset: int, value: int) -> None:
            assert offset in {0xF8, 0xFC}

        def write_config_u8(self, offset: int, value: int) -> None:
            assert offset == 0xE0
            self.control_writes.append(value)

    controller = FakeController()

    data = _read_memory(controller, mem_type=0x20, size=3, addr=0, sec_low=0x9F)

    assert data == bytes((0xEF, 0x30, 0x11))
    assert controller.control_writes == [0x02, 0x02, 0x01, 0x01]


def _build_2214a_image(raw_payload: bytes) -> bytes:
    header_len = 0x30
    payload_offset = header_len + 9
    total_size = payload_offset + len(raw_payload)
    if total_size % 2:
        total_size += 1

    data = bytearray(total_size)
    data[4:6] = struct.pack("<H", header_len)
    header_tags = b"2214A_RCFG\x002214A_FW\x00"
    data[8 : 8 + len(header_tags)] = header_tags
    data[header_len + 5 : header_len + 9] = struct.pack("<I", len(raw_payload))
    data[payload_offset : payload_offset + len(raw_payload)] = raw_payload
    return bytes(data)
