"""PeMMU: per-PE virtual-to-physical address translation.

Page-aligned dict lookup for O(1) VA→PA translation.
Used as a utility class by PE_DMA, PE_GEMM (direct call),
and as a component inbox target for MmuMapMsg/MmuUnmapMsg.
"""
from __future__ import annotations


class PageFault(Exception):
    """Raised when VA has no mapping in the page table."""

    def __init__(self, va: int | str) -> None:
        if isinstance(va, str):
            super().__init__(va)
        else:
            self.va = va
            super().__init__(f"PageFault at VA 0x{va:x}")


class PeMMU:
    """Per-PE MMU with sub-page-capable VA→PA translation table.

    Each page-table entry is a list of (start_in_page, end_in_page,
    pa_at_offset_zero) regions. This is a SIMULATOR STOPGAP — real MMUs
    store one PA per page-table entry. Sub-page regions exist here so
    DPPolicy layouts that shard below page granularity (e.g. 128 B
    payloads with 4 KB pages) don't silently mis-route through last-
    write-wins overwrites. Memory note: project_mmu_subpage_stopgap.md.

    Args:
        page_size: Page size in bytes (default 2 MB).
        overhead_ns: Per-access TLB lookup latency in nanoseconds.
    """

    def __init__(
        self,
        page_size: int = 2 * 1024 * 1024,
        overhead_ns: float = 0.0,
    ) -> None:
        self._page_size = page_size
        self._page_shift = (page_size - 1).bit_length()
        self._page_mask = page_size - 1
        # vpn → list of (start_in_page, end_in_page, pa_at_offset_zero).
        # pa_at_offset_zero is the PA that offset 0 of the page would map
        # to under this region — i.e. translate(off) = pa_at_offset_zero
        # + off when start <= off < end.
        self._table: dict[int, list[tuple[int, int, int]]] = {}
        self._overhead_ns = overhead_ns

    @property
    def overhead_ns(self) -> float:
        return self._overhead_ns

    @property
    def num_entries(self) -> int:
        return len(self._table)

    def map(self, va: int, pa: int, size: int) -> None:
        """Register VA→PA mapping for a contiguous range.

        Sub-page-aware: a single page can hold multiple disjoint regions,
        each pointing to a different PA. Later map() calls APPEND a new
        region; on overlap with an existing region, the new region wins
        for the overlapping offsets (translate iterates in reverse so the
        last write takes precedence — matches legacy single-PA behavior
        when a full page is re-mapped).
        """
        end_va = va + size
        cur = va
        while cur < end_va:
            vpn = cur >> self._page_shift
            page_base_va = vpn << self._page_shift
            page_end_va = page_base_va + self._page_size
            region_start = cur - page_base_va
            region_end = min(end_va, page_end_va) - page_base_va
            # PA seen at offset 0 of page if this region's mapping covered it
            pa_at_offset_zero = pa + (cur - va) - region_start
            self._table.setdefault(vpn, []).append(
                (region_start, region_end, pa_at_offset_zero)
            )
            cur = page_base_va + region_end

    def unmap(self, va: int, size: int) -> None:
        """Remove VA mapping for a contiguous range.

        Drops any region whose extent is contained within the unmapped
        range. Partial overlaps (region straddles the range boundary)
        are left in place — caller is expected to unmap on the same
        boundaries it mapped on.
        """
        end_va = va + size
        cur = va
        while cur < end_va:
            vpn = cur >> self._page_shift
            page_base_va = vpn << self._page_shift
            page_end_va = page_base_va + self._page_size
            unmap_start = cur - page_base_va
            unmap_end = min(end_va, page_end_va) - page_base_va
            regions = self._table.get(vpn)
            if regions is not None:
                kept = [
                    r for r in regions
                    if not (r[0] >= unmap_start and r[1] <= unmap_end)
                ]
                if kept:
                    self._table[vpn] = kept
                else:
                    del self._table[vpn]
            cur = page_base_va + unmap_end

    def translate(self, va: int) -> int:
        """Translate VA to PA. Raises PageFault if unmapped."""
        vpn = va >> self._page_shift
        regions = self._table.get(vpn)
        if regions is None:
            raise PageFault(va)
        offset = va & self._page_mask
        # Iterate latest-first so newer map() calls win on overlap
        for start, end, pa_at_offset_zero in reversed(regions):
            if start <= offset < end:
                return pa_at_offset_zero + offset
        raise PageFault(va)
