"""KernelRunner: greenlet-based kernel ↔ SimPy bridge (ADR-0020 D3).

Replaces Phase 0 (static command list) with interleaved execution:
  - tl.load()  → SimPy DMA timing + MemoryStore read → real data to kernel
  - tl.store() → MemoryStore write + SimPy DMA timing
  - tl.composite(gemm) → SimPy timing + op_log (actual compute in Phase 2)

The kernel runs as a child greenlet; SimPy loop is the parent.
"""
from __future__ import annotations

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

import simpy
from greenlet import greenlet

from kernbench.common.ipcq_types import IpcqRecvCmd, IpcqRequest, IpcqSendCmd, RecvFuture
from kernbench.common.pe_commands import (
    CompletionHandle,
    CompositeCmd,
    DmaReadCmd,
    DmaWriteCmd,
    GemmCmd,
    MathCmd,
    PeCommand,
    PeCpuOverheadCmd,
    PeInternalTxn,
    TensorHandle,
    WaitCmd,
)

if TYPE_CHECKING:
    from kernbench.sim_engine.memory_store import MemoryStore


class KernelRunner:
    """Greenlet ↔ SimPy bridge for kernel execution (ADR-0020 D3).

    PE_CPU creates a KernelRunner and yields from its run() method.
    The kernel function executes as a plain Python function inside a
    child greenlet, using TLContext methods that switch back to SimPy.
    """

    def __init__(
        self,
        pe_prefix: str,
        pe_idx: int,
        sip_idx: int,
        cube_idx: int,
        scheduler_id: str,
        out_ports: dict[str, simpy.Store],
        store: MemoryStore | None = None,
        num_cubes: int = 1,
        ipcq_id: str | None = None,
        scratch_base: int = 0,
        scratch_size: int = 1 << 20,
        cost_model: Any = None,
        clock_freq_ghz: float = 1.0,
    ) -> None:
        self._pe_prefix = pe_prefix
        self._pe_idx = pe_idx
        self._sip_idx = sip_idx
        self._cube_idx = cube_idx
        self._num_cubes = num_cubes
        self._scheduler_id = scheduler_id
        self._out_ports = out_ports
        self._store = store
        self._parent: greenlet | None = None
        # Optional IPCQ port (ADR-0023). If None, IPCQ commands raise.
        self._ipcq_id = ipcq_id or f"{pe_prefix}.pe_ipcq"
        # PE-local scratch for compute output TensorHandles (ADR-0020 D3
        # extension). The TLContext allocates from this pool when math/dot
        # ops produce a result that may later be used as a send/store source.
        self._scratch_base = scratch_base
        self._scratch_size = scratch_size
        # ADR-0064 Rev2 structural dispatch cost model (+ clock for cycle→ns).
        self._cost_model = cost_model
        self._clock_freq_ghz = clock_freq_ghz

    def run(
        self,
        env: simpy.Environment,
        kernel_fn: Any,
        kernel_args: list,
        num_programs: int,
    ) -> Generator:
        """SimPy generator: run kernel with greenlet interleaving.

        This is the SimPy-side loop. It:
        1. Creates a TLContext connected to this runner
        2. Spawns the kernel in a child greenlet
        3. Receives commands via greenlet.switch
        4. Dispatches each command through SimPy components
        5. Returns results to the kernel
        """
        from kernbench.triton_emu.tl_context import TLContext

        self._parent = greenlet.getcurrent()

        tl = TLContext(
            pe_id=self._pe_idx,
            num_programs=num_programs,
            cube_id=self._cube_idx,
            num_cubes=self._num_cubes,
            dispatch_cycles=0,
            runner=self,
            scratch_base=self._scratch_base,
            scratch_size=self._scratch_size,
            cost_model=self._cost_model,
            clock_freq_ghz=self._clock_freq_ghz,
        )
        self._tl = tl  # exposed so switch_to_simpy can re-set on restore

        def _kernel_entry():
            TLContext._set_active(tl)  # type: ignore[attr-defined]
            try:
                kernel_fn(*kernel_args, tl=tl)
            finally:
                TLContext._set_active(None)  # type: ignore[attr-defined]
            return None  # signal kernel completion

        g = greenlet(_kernel_entry)
        pending: dict[str, simpy.Event] = {}
        composite_results: list[dict] = []

        # Helper: set our tl as active just before resuming the kernel.
        # Multiple PE kernel runners share the same thread-local; without
        # this, another runner's kernel may have left a different context.
        def _switch_kernel(*args):
            TLContext._set_active(tl)  # type: ignore[attr-defined]
            return g.switch(*args)

        # Start kernel — first switch returns first command (or None if kernel is done)
        cmd = _switch_kernel()

        while cmd is not None:
            if isinstance(cmd, PeCpuOverheadCmd):
                yield env.timeout(cmd.cycles)
                cmd = _switch_kernel()

            elif isinstance(cmd, WaitCmd):
                if cmd.handle is not None:
                    evt = pending.pop(cmd.handle.id, None)
                    if evt:
                        yield evt
                else:
                    for evt in pending.values():
                        yield evt
                    pending.clear()
                cmd = _switch_kernel()

            elif isinstance(cmd, DmaReadCmd):
                # Legacy blocking path — retained as a fallback for any
                # caller that bypasses the lazy ``tl.load`` surface. New
                # lazy loads come in as ("load_issue", future) below.
                done_evt = env.event()
                pe_txn = PeInternalTxn(
                    command=cmd, done=done_evt, pe_prefix=self._pe_prefix,
                )
                yield self._out_ports[self._scheduler_id].put(pe_txn)
                yield done_evt

                # Read actual data from MemoryStore (if available)
                data = None
                if self._store is not None:
                    try:
                        data = self._store.read(
                            "hbm", cmd.src_addr,
                            shape=cmd.handle.shape, dtype=cmd.handle.dtype,
                        )
                    except KeyError:
                        pass
                cmd = _switch_kernel(data)

            elif isinstance(cmd, DmaWriteCmd):
                # Write to MemoryStore first (visibility = issue, ADR-0020 D3).
                # When data is None (e.g. timing-only TensorHandle math result),
                # this is a no-op; Phase 2 dma_write replay handles those.
                if self._store is not None and cmd.handle.data is not None:
                    self._store.write("hbm", cmd.dst_addr, cmd.handle.data)

                done_evt = env.event()
                pe_txn = PeInternalTxn(
                    command=cmd, done=done_evt, pe_prefix=self._pe_prefix,
                )
                yield self._out_ports[self._scheduler_id].put(pe_txn)
                yield done_evt
                cmd = _switch_kernel()

            elif isinstance(cmd, CompositeCmd):
                # Non-blocking composite
                done_evt = env.event()
                pe_txn = PeInternalTxn(
                    command=cmd, done=done_evt, pe_prefix=self._pe_prefix,
                )
                composite_results.append(pe_txn.result_data)
                yield self._out_ports[self._scheduler_id].put(pe_txn)
                pending[cmd.completion.id] = done_evt
                cmd = _switch_kernel()

            elif isinstance(cmd, (GemmCmd, MathCmd)):
                # Blocking compute command
                done_evt = env.event()
                pe_txn = PeInternalTxn(
                    command=cmd, done=done_evt, pe_prefix=self._pe_prefix,
                )
                yield self._out_ports[self._scheduler_id].put(pe_txn)
                yield done_evt
                cmd = _switch_kernel()

            elif isinstance(cmd, IpcqSendCmd):
                # Forward IpcqRequest to PE_IPCQ, wait for done
                if self._ipcq_id not in self._out_ports:
                    raise RuntimeError(
                        f"PE_IPCQ port {self._ipcq_id!r} not wired to runner"
                    )
                done_evt = env.event()
                req = IpcqRequest(command=cmd, done=done_evt)
                yield self._out_ports[self._ipcq_id].put(req)
                yield done_evt
                cmd = _switch_kernel()

            elif isinstance(cmd, IpcqRecvCmd):
                if self._ipcq_id not in self._out_ports:
                    raise RuntimeError(
                        f"PE_IPCQ port {self._ipcq_id!r} not wired to runner"
                    )
                done_evt = env.event()
                req = IpcqRequest(command=cmd, done=done_evt)
                yield self._out_ports[self._ipcq_id].put(req)
                yield done_evt
                # Read actual data from MemoryStore at the slot address
                data = None
                src_space = req.result_data.get("src_space", "tcm")
                src_addr = req.result_data.get("src_addr", 0)
                if self._store is not None:
                    try:
                        data = self._store.read(
                            src_space, src_addr,
                            shape=cmd.shape, dtype=cmd.dtype,
                        )
                    except KeyError:
                        pass
                # Build result dict for tl.recv to wrap in TensorHandle
                result = {
                    "data": data,
                    "src_space": src_space,
                    "src_addr": src_addr,
                    "direction": req.result_data.get("direction", cmd.direction),
                    "dtype": cmd.dtype,
                    "shape": cmd.shape,
                    "nbytes": req.result_data.get("nbytes", 0),
                }
                cmd = _switch_kernel(result)

            elif isinstance(cmd, tuple) and len(cmd) == 2 and cmd[0] == "load_issue":
                # ADR-0062 §D2: lazy tl.load. Post the DmaReadCmd, store the
                # done event on the LoadFuture, switch back immediately — do
                # NOT yield done_evt here. The auto-wait at first use
                # ("load_await" below) is what eventually yields it.
                future = cmd[1]
                done_evt = env.event()
                pe_txn = PeInternalTxn(
                    command=future.cmd, done=done_evt, pe_prefix=self._pe_prefix,
                )
                yield self._out_ports[self._scheduler_id].put(pe_txn)
                future.event = done_evt
                cmd = _switch_kernel(None)

            elif isinstance(cmd, tuple) and len(cmd) == 2 and cmd[0] == "load_await":
                # ADR-0062 §D2: auto-wait at first use. Yield on the future's
                # DMA event if not yet triggered, then read data and attach
                # it to the handle (frozen dataclass — mutate via object.
                # __setattr__, the same controlled pattern Phase 1
                # blocking tl.load used).
                handle = cmd[1]
                future = handle.pending
                if future is not None and not future.resolved:
                    if future.event is not None and not future.event.triggered:
                        yield future.event
                    data = None
                    if self._store is not None:
                        try:
                            data = self._store.read(
                                "hbm", future.cmd.src_addr,
                                shape=handle.shape, dtype=handle.dtype,
                            )
                        except KeyError:
                            pass
                    future.data = data
                    future.resolved = True
                    object.__setattr__(handle, "data", data)
                cmd = _switch_kernel(None)

            elif isinstance(cmd, tuple) and len(cmd) == 2 and cmd[0] == "recv_async":
                # Non-blocking recv: post the IpcqRequest now, store the
                # event in the future, return None to kernel.
                future: RecvFuture = cmd[1]
                done_evt = env.event()
                req = IpcqRequest(command=future.cmd, done=done_evt)
                future.request = req
                future.event = done_evt
                yield self._out_ports[self._ipcq_id].put(req)
                cmd = _switch_kernel(None)

            elif isinstance(cmd, tuple) and len(cmd) == 2 and cmd[0] == "recv_wait":
                future = cmd[1]
                if not future.event.triggered:
                    yield future.event
                req = future.request
                src_space = req.result_data.get("src_space", "tcm")
                src_addr = req.result_data.get("src_addr", 0)
                data = None
                if self._store is not None:
                    try:
                        data = self._store.read(
                            src_space, src_addr,
                            shape=future.cmd.shape, dtype=future.cmd.dtype,
                        )
                    except KeyError:
                        pass
                result = {
                    "data": data,
                    "src_space": src_space,
                    "src_addr": src_addr,
                    "direction": req.result_data.get("direction", future.cmd.direction),
                    "dtype": future.cmd.dtype,
                    "shape": future.cmd.shape,
                    "nbytes": req.result_data.get("nbytes", 0),
                }
                cmd = _switch_kernel(result)

            else:
                # Unknown command — pass through as blocking
                done_evt = env.event()
                pe_txn = PeInternalTxn(
                    command=cmd, done=done_evt, pe_prefix=self._pe_prefix,
                )
                yield self._out_ports[self._scheduler_id].put(pe_txn)
                yield done_evt
                cmd = _switch_kernel()

        # Wait remaining pending composites
        for evt in pending.values():
            yield evt

        # Return composite results for PE_CPU aggregation
        self._composite_results = composite_results

    def switch_to_simpy(self, cmd: PeCommand) -> Any:
        """Called from TLContext (child greenlet) to send command to SimPy.

        Returns the result from SimPy (e.g., numpy array for DMA read).
        """
        assert self._parent is not None
        return self._parent.switch(cmd)
