"""Op log infrastructure for 2-pass data execution (ADR-0020 D2, D5).

OpRecord: single data operation with timing, params, and dependencies.
OpLogger: collects OpRecords from ComponentBase hooks during Phase 1.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any


@dataclass
class OpRecord:
    """Single data operation record (ADR-0020 D5)."""

    t_start: float
    t_end: float
    component_id: str
    op_kind: str               # "memory" | "gemm" | "math"
    op_name: str               # e.g. "dma_read", "gemm_f16", "exp"
    params: dict[str, Any]
    dependency_ids: list[int] = field(default_factory=list)


class OpLogger:
    """Collects OpRecords during Phase 1 simulation (ADR-0020 D2).

    Thread-safe is not required — SimPy is single-threaded.
    Records are maintained in t_start stable ordering (insertion order).
    """

    def __init__(self, memory_store: Any | None = None) -> None:
        self._records: list[OpRecord] = []
        self._pending: dict[int, dict[str, Any]] = {}  # msg id → partial record
        # Optional MemoryStore reference. When set, math op records capture
        # input data snapshots at record_end time so Phase 2 replay does
        # not depend on slot/scratch addrs surviving until math runs.
        self._memory_store = memory_store

    @property
    def records(self) -> list[OpRecord]:
        """Records sorted by t_start (stable ordering per ADR-0020 D5)."""
        self._records.sort(key=lambda r: r.t_start)
        return self._records

    def record_start(self, t: float, component_id: str, msg: Any) -> None:
        """Called by ComponentBase._on_process_start.

        Snapshots TileToken stage_type at start time so we can attribute the
        record correctly even if the token advances stage_idx before
        record_end fires.
        """
        snap: dict[str, Any] = {}
        # TileToken (ADR-0014 D6 pipeline) — capture which stage this is and its
        # per-stage params (e.g. op_kind/scope for epilogue MATH stages) so
        # we can recover them at record_end even after the token advances.
        try:
            stage = getattr(msg, "current_stage", None)
            if stage is not None:
                snap["stage_type"] = stage.stage_type.name
                if isinstance(getattr(stage, "params", None), dict):
                    snap["stage_params"] = dict(stage.params)
        except Exception:
            pass
        self._pending[id(msg)] = {
            "t_start": t,
            "component_id": component_id,
            "msg": msg,
            "snap": snap,
        }

    def record_end(self, t: float, component_id: str, msg: Any) -> None:
        """Called by ComponentBase._on_process_end."""
        pending = self._pending.pop(id(msg), None)
        if pending is None:
            return
        op_kind, op_name, params = _extract_op_info(msg)
        # Merge TileToken stage_type captured at record_start into params,
        # and reflect it in op_name so reporting can disambiguate
        # DMA_READ vs DMA_WRITE and FETCH vs STORE on the same component.
        snap = pending.get("snap", {})
        stage_type = snap.get("stage_type")
        if stage_type is not None:
            params = dict(params)
            # Merge per-stage params (e.g. op_kind, scope) captured at start.
            stage_params = snap.get("stage_params")
            if isinstance(stage_params, dict):
                for k, v in stage_params.items():
                    params.setdefault(k, v)
            params["stage_type"] = stage_type
            if op_name == "TileToken":
                op_name = f"TileToken/{stage_type}"
            # ADR-0065 N2: a recipe MATH stage carries operand addresses →
            # promote it to a computable math op so the DataExecutor runs it.
            # Legacy epilogue MATH stages (no addresses) stay op_kind-opaque
            # (latency-only) → byte-equal.
            if stage_type == "MATH" and "input_addrs" in params:
                op_kind = "math"
        # Snapshot data at record time so Phase 2 replay sidesteps
        # downstream mutations of source addrs (e.g. a tl.store that
        # overwrites HBM after a load handle was sent, or a slot that
        # gets reused on the next ring round).
        if self._memory_store is not None:
            if op_kind == "math":
                handle_snaps = params.get("input_handle_data") or ()
                snaps: list[Any] = []
                for i, (addr, shape, space, idtype) in enumerate(zip(
                    params.get("input_addrs", []),
                    params.get("input_shapes", []),
                    params.get("input_spaces", []),
                    params.get("input_dtypes", []),
                )):
                    if i < len(handle_snaps) and handle_snaps[i] is not None:
                        carried = handle_snaps[i]
                        snaps.append(
                            carried.copy() if hasattr(carried, "copy") else carried
                        )
                        continue
                    try:
                        arr = self._memory_store.read(
                            space, addr, shape=shape, dtype=idtype,
                        )
                        snaps.append(arr.copy() if hasattr(arr, "copy") else arr)
                    except Exception:
                        snaps.append(None)
                params["input_snapshots"] = snaps
                params.pop("input_handle_data", None)
            elif op_name == "dma_write":
                # ADR-0027 fix: only snapshot HBM sources. TCM (PE scratch)
                # sources are repopulated by Phase 2 math/gemm replay —
                # capturing a Phase-1-time snapshot here would pick up stale
                # data from a PRIOR kernel's Phase 2 output that aliased the
                # same scratch address, causing the later kernel's replay
                # to write that stale value instead of the fresh math
                # result. See ADR-0027 postmortem (TP gemm → all_reduce).
                if params.get("src_space") == "hbm":
                    try:
                        arr = self._memory_store.read(
                            params["src_space"], params["src_addr"],
                            shape=params.get("shape"), dtype=params.get("dtype"),
                        )
                        params["snapshot"] = (
                            arr.copy() if hasattr(arr, "copy") else arr
                        )
                    except Exception:
                        params["snapshot"] = None
        self._records.append(OpRecord(
            t_start=pending["t_start"],
            t_end=t,
            component_id=pending["component_id"],
            op_kind=op_kind,
            op_name=op_name,
            params=params,
        ))

    def record_copy(
        self, t_start: float, t_end: float, component_id: str,
        src_space: str, src_addr: int,
        dst_space: str, dst_addr: int,
        shape: tuple[int, ...], dtype: str, nbytes: int,
        snapshot: Any = None,
    ) -> None:
        """Record a memory copy op for Phase 2 replay (ADR-0023 + ADR-0020).

        ``snapshot``: if provided (e.g. token.data from in-flight DMA),
        used directly. Otherwise falls back to a fresh read from
        MemoryStore[src_addr]. The snapshot is what Phase 2 writes into
        dst_addr, avoiding stale-source races from cross-PE mutations.
        """
        snap = snapshot
        if snap is None and self._memory_store is not None:
            try:
                arr = self._memory_store.read(
                    src_space, src_addr, shape=shape, dtype=dtype,
                )
                snap = arr.copy() if hasattr(arr, "copy") else arr
            except Exception:
                snap = None
        self._records.append(OpRecord(
            t_start=t_start, t_end=t_end,
            component_id=component_id,
            op_kind="memory", op_name="ipcq_copy",
            params={
                "src_space": src_space, "src_addr": src_addr,
                "dst_space": dst_space, "dst_addr": dst_addr,
                "shape": shape, "dtype": dtype, "nbytes": nbytes,
                "snapshot": snap,
            },
        ))


def _extract_op_info(msg: Any) -> tuple[str, str, dict[str, Any]]:
    """Extract op_kind, op_name, params from a data_op message."""
    from kernbench.common.pe_commands import (
        CompositeCmd, CopyCmd, DmaReadCmd, DmaWriteCmd, GemmCmd, MathCmd,
    )
    if isinstance(msg, DmaReadCmd):
        return "memory", "dma_read", {
            "src_addr": msg.src_addr,
            "nbytes": msg.nbytes,
            "handle_id": msg.handle.id,
        }
    if isinstance(msg, DmaWriteCmd):
        return "memory", "dma_write", {
            "src_space": getattr(msg.handle, "space", "tcm"),
            "src_addr": msg.handle.addr,
            "shape": msg.handle.shape,
            "dtype": msg.handle.dtype,
            "dst_space": "hbm",
            "dst_addr": msg.dst_addr,
            "nbytes": msg.nbytes,
            "handle_id": msg.handle.id,
        }
    if isinstance(msg, GemmCmd):
        return "gemm", f"gemm_{msg.a.dtype}", {
            "src_a_addr": msg.a.addr,
            "src_b_addr": msg.b.addr,
            "dst_addr": msg.out.addr,
            "shape_a": msg.a.shape,
            "shape_b": msg.b.shape,
            "shape_out": msg.out.shape,
            "dtype_in": msg.a.dtype,
            "dtype_out": msg.out.dtype,
            "m": msg.m, "k": msg.k, "n": msg.n,
            # ADR-0027: preserve per-operand + output MemoryStore spaces so
            # Phase 2 replay can resolve HBM-resident operands (e.g. tl.load
            # results keep space="hbm"). Absent → DataExecutor falls back
            # to the legacy single-space mode via ``addr_space``.
            "src_a_space": getattr(msg.a, "space", "tcm"),
            "src_b_space": getattr(msg.b, "space", "tcm"),
            "dst_space": getattr(msg.out, "space", "tcm"),
        }
    if isinstance(msg, MathCmd):
        return "math", msg.op, {
            "input_addrs": [h.addr for h in msg.inputs],
            "input_shapes": [h.shape for h in msg.inputs],
            "input_spaces": [getattr(h, "space", "tcm") for h in msg.inputs],
            "input_dtypes": [h.dtype for h in msg.inputs],
            "input_handle_data": tuple(getattr(h, "data", None) for h in msg.inputs),
            "dst_addr": msg.out.addr,
            "dst_space": getattr(msg.out, "space", "tcm"),
            "shape_out": msg.out.shape,
            "dtype": msg.out.dtype,
            "axis": msg.axis,
        }
    if isinstance(msg, CopyCmd):
        return "math", "copy", {
            "src_addr": msg.src.addr,
            "src_space": getattr(msg.src, "space", "tcm"),
            "dst_addr": msg.dst.addr,
            "dst_space": getattr(msg.dst, "space", "tcm"),
            "shape": msg.src.shape,
            "dtype": msg.src.dtype,
            "nbytes": msg.nbytes,
        }
    if isinstance(msg, CompositeCmd):
        # Flat-ops (ADR-0065 D1/D5): scan for the GEMM op — it is the head
        # for a legacy composite (index 0) but sits AFTER the prologue MATH
        # ops in a recipe composite. Its a/b/out give the matmul the Phase 2
        # DataExecutor replays as a single GemmCmd-equivalent.
        gemm = next((o for o in msg.ops if o.kind == "gemm"), None)
        if gemm is not None:
            a = gemm.operands.get("a")
            b = gemm.operands.get("b")
            # Accumulate (O += a@b) when an `add` epilogue targets the GEMM's
            # own output (the recipe pattern: GEMM P·V folds into the already-
            # rescaled O). Default overwrite for a plain GEMM composite.
            accumulate = any(
                o.kind == "add" and o.operands.get("other") is gemm.out
                for o in msg.ops
            )
            params: dict[str, Any] = {
                "op": "gemm",
                "out_addr": gemm.out.addr,
                "out_nbytes": gemm.out.nbytes,
                "accumulate": accumulate,
            }
            if a is not None and b is not None:
                params.update({
                    "src_a_addr": a.addr,
                    "src_b_addr": b.addr,
                    "shape_a": a.shape,
                    "shape_b": b.shape,
                    "dtype_in": a.dtype,
                    "dtype_out": gemm.out.dtype,
                    "src_a_space": getattr(a, "space", "tcm"),
                    "src_b_space": getattr(b, "space", "tcm"),
                    "dst_space": getattr(gemm.out, "space", "tcm"),
                    "dst_addr": gemm.out.addr,
                })
            return "gemm", "composite_gemm", params
        # MATH-only composite (no GEMM) — fall back to the head op.
        head = msg.ops[0]
        return "math", f"composite_{head.kind}", {
            "op": head.kind,
            "out_addr": head.out.addr,
            "out_nbytes": head.out.nbytes,
        }
    # Fallback for unknown data_op messages
    return "unknown", type(msg).__name__, {}
