"""PE_MMU component: address translation unit.

Component role: receives MmuMapMsg/MmuUnmapMsg via inbox (independent of PE_CPU).
Utility role: PE_DMA/PE_GEMM call mmu.translate() directly (no SimPy overhead).
"""
from __future__ import annotations

from collections.abc import Generator
from typing import TYPE_CHECKING, Any

import simpy

from kernbench.components.base import ComponentBase, ComponentRegistry
from kernbench.policy.address.pe_mmu import PeMMU

if TYPE_CHECKING:
    from kernbench.components.context import ComponentContext
    from kernbench.topology.types import Node


class PeMmuComponent(ComponentBase):
    """PE_MMU: per-PE virtual-to-physical address translation.

    Receives MmuMapMsg/MmuUnmapMsg via inbox and updates the internal
    page table. PE_DMA and PE_GEMM access the underlying PeMMU object
    via the ``mmu`` property for synchronous VA→PA translation.
    """

    def __init__(self, node: Node, ctx: ComponentContext | None = None) -> None:
        super().__init__(node, ctx)
        page_size = int(node.attrs.get("page_size", 2 * 1024 * 1024))
        overhead_ns = float(node.attrs.get("tlb_overhead_ns", 0.0))
        self._mmu = PeMMU(page_size=page_size, overhead_ns=overhead_ns)

    @property
    def mmu(self) -> PeMMU:
        """The underlying PeMMU utility object for direct translate() calls."""
        return self._mmu

    def run(self, env: simpy.Environment, nbytes: int) -> Generator:
        yield env.timeout(0)

    def _worker(self, env: simpy.Environment) -> Generator:
        """Process MmuMapMsg/MmuUnmapMsg from inbox."""
        from kernbench.runtime_api.kernel import MmuMapMsg, MmuUnmapMsg

        while True:
            txn: Any = yield self._inbox.get()

            if hasattr(txn, "request"):
                request = txn.request
                if isinstance(request, MmuMapMsg):
                    for entry in request.entries:
                        self._mmu.map(
                            va=entry["va"], pa=entry["pa"], size=entry["size"],
                        )
                    txn.done.succeed()
                elif isinstance(request, MmuUnmapMsg):
                    for entry in request.entries:
                        self._mmu.unmap(va=entry["va"], size=entry["size"])
                    txn.done.succeed()
                else:
                    # Forward non-MMU transactions normally
                    yield from self._forward_txn(env, txn)
            else:
                yield from self._forward_txn(env, txn)