"""WPR2 / FWSEC state registers.

Evidence kind: **hardware register/field header + firmware use** (the boot
flow locates the firmware write-protected region from these registers on
pre-Hopper parts; the host/booter writes them, the firmware reads them).

Modelled registers (offsets in the FB/PFB register aperture):

* ``NV_PFB_PRI_MMU_WPR2_ADDR_LO`` ``0x1FA824``: WPR2 start, bits ``31:4``
  shifted by ``NV_PFB_PRI_MMU_WPR2_ADDR_LO_ALIGNMENT = 12`` (4 KiB).
* ``NV_PFB_PRI_MMU_WPR2_ADDR_HI`` ``0x1FA828``: points at the start of the
  last 128 KiB block of the WPR, same alignment; the end is
  ``(value << 12) + 0x20000``.

A profile that wants the firmware to see a WPR declares ``wpr2_start`` /
``wpr2_end``; without a declaration the window is unhandled and logged.
FWSEC-specific state is not modelled (no proven register use yet).
"""
from __future__ import annotations

from .base import Device

WPR2_ADDR_LO = 0x1FA824
WPR2_ADDR_HI = 0x1FA828
WPR2_ALIGNMENT = 12
WPR2_BLOCK = 0x20000


class Wpr2Block(Device):
    """WPR2 address pair as seen by the firmware (host/booter-written)."""

    name = "fwsec"

    def __init__(self, *, base: int = 0, size: int = 0x1000,
                 wpr2_start: int | None = None, wpr2_end: int | None = None):
        super().__init__(base=base, size=size)
        self.reset()
        if wpr2_start is not None and wpr2_end is not None:
            self.set_range(wpr2_start, wpr2_end)

    def set_range(self, start: int, end: int) -> None:
        if end <= start:
            raise ValueError("WPR2 end must be greater than start")
        self.regs[WPR2_ADDR_LO] = (start >> WPR2_ALIGNMENT) & 0xFFFFFFFF
        self.regs[WPR2_ADDR_HI] = ((end - WPR2_BLOCK) >> WPR2_ALIGNMENT) & 0xFFFFFFFF

    @property
    def start(self) -> int | None:
        """WPR2 start derived from ADDR_LO (``None`` when unset)."""
        lo = self.regs.get(WPR2_ADDR_LO)
        return None if lo is None else lo << WPR2_ALIGNMENT

    @property
    def end(self) -> int | None:
        """WPR2 end derived from ADDR_HI (``None`` when unset)."""
        hi = self.regs.get(WPR2_ADDR_HI)
        return None if hi is None else (hi << WPR2_ALIGNMENT) + WPR2_BLOCK

    def reset(self) -> None:
        self.regs: dict[int, int] = {}

    def read(self, offset: int, size: int) -> int | None:
        if size == 4 and offset in self.regs:
            return self.regs[offset]
        return None

    def write(self, offset: int, size: int, value: int) -> bool:
        if size == 4 and offset in (WPR2_ADDR_LO, WPR2_ADDR_HI):
            self.regs[offset] = value & 0xFFFFFFFF
            return True
        return False

    def state(self) -> dict:
        return {"regs": {str(k): v for k, v in sorted(self.regs.items())},
                "start": self.start, "end": self.end}

    def load_state(self, data: dict) -> None:
        # start/end are derived from regs; serialised copies are ignored so a
        # restored block can never report a stale range.
        self.regs = {int(k): int(v) for k, v in data.get("regs", {}).items()}