"""PE_SCHEDULER: plan generation + tile dispatch (ADR-0014 D6).

Receives PeInternalTxn from PE_CPU, routes to engines:
  - Simple commands (DmaReadCmd, GemmCmd, etc.) → direct dispatch to engine
  - CompositeCmd → generate TilePlan, feed tiles via _feed_loop

Composite pipeline uses token self-routing (ADR-0014 D6):
  Scheduler only does initial dispatch + completion tracking.
  Tiles chain through components based on their plan's stage sequence.
"""
from __future__ import annotations

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

import simpy

from kernbench.components.base import ComponentBase

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


class _RwHazardTracker:
    """Strict-FIFO cross-composite RW hazard tracker (ADR-0065 D6.3 / DDD §8).

    A composite is registered as in-flight by its *write* set
    (``rw_handles``). A new composite whose read **or** write handles
    intersect any in-flight write set waits on those composites' done events
    before being admitted. Composites with no ``rw_handles`` (legacy) are
    never registered and never block — existing benches are unaffected.

    Strict FIFO is conservative: even partial overlap waits for the prior
    overlapping composite to fully drain (RW-aware reorder is ADR-0065 A4,
    deferred).
    """

    def __init__(self) -> None:
        # (completion_id, write-handle id set, done_event)
        self._inflight: list[tuple[str, set, Any]] = []

    @staticmethod
    def _conflict_ids(cmd: Any) -> set:
        """Handle ids this composite reads (op operands) or writes (rw)."""
        ids = {h.id for h in cmd.rw_handles}
        for op in cmd.ops:
            for h in op.operands.values():
                hid = getattr(h, "id", None)
                if hid is not None:
                    ids.add(hid)
        return ids

    def admit(self, env: Any, cmd: Any, done_event: Any) -> Generator:
        """Block until no in-flight composite's write set overlaps this
        composite's read/write set, then register this composite's writes."""
        cids = self._conflict_ids(cmd)
        while True:
            blocking = [ev for _cid, rw, ev in self._inflight if rw & cids]
            if not blocking:
                break
            yield blocking[0]
        if cmd.rw_handles:
            self._inflight.append(
                (cmd.completion.id, {h.id for h in cmd.rw_handles}, done_event)
            )

    def retire(self, completion_id: str) -> None:
        self._inflight = [t for t in self._inflight if t[0] != completion_id]


class PeSchedulerComponent(ComponentBase):
    """PE_SCHEDULER: sole dispatcher inside a PE (ADR-0014 D1, D6).

    Simple commands are forwarded to the appropriate engine.
    CompositeCmd creates a TilePlan and feeds tiles into the pipeline.

    Single _feed_loop process per scheduler ensures FIFO command ordering.
    """

    TILE_M = 32
    TILE_K = 64
    TILE_N = 32

    _CMD_DISPATCH: dict[type, str] = {}

    @classmethod
    def _ensure_dispatch_table(cls) -> None:
        if cls._CMD_DISPATCH:
            return
        from kernbench.common.pe_commands import (
            CopyCmd, DmaReadCmd, DmaWriteCmd, GemmCmd, MathCmd,
        )
        cls._CMD_DISPATCH = {
            DmaReadCmd: "pe_dma",
            DmaWriteCmd: "pe_dma",
            GemmCmd: "pe_gemm",
            MathCmd: "pe_math",
            # ADR-0063 §D3.1: tl.copy_to → vector engine.
            CopyCmd: "pe_math",
        }

    def __init__(self, node: Node, ctx: ComponentContext | None = None) -> None:
        super().__init__(node, ctx)
        self._pe_prefix = node.id.rsplit(".", 1)[0]
        self._ensure_dispatch_table()
        self._pending_feeds: simpy.Store | None = None
        self._pipeline_counter = 0
        self._hazard = _RwHazardTracker()  # ADR-0065 D6.3 strict-FIFO RW

    def start(self, env: simpy.Environment) -> None:
        self._pending_feeds = simpy.Store(env)
        super().start(env)
        env.process(self._feed_loop(env))

    def run(self, env: simpy.Environment, nbytes: int) -> Generator:
        overhead_ns = float(self.node.attrs.get("overhead_ns", 0.0))
        yield env.timeout(overhead_ns)

    def _worker(self, env: simpy.Environment) -> Generator:
        from kernbench.common.pe_commands import PeInternalTxn

        while True:
            msg: Any = yield self._inbox.get()
            if isinstance(msg, PeInternalTxn):
                env.process(self._dispatch(env, msg))
            else:
                yield from self._forward_txn(env, msg)

    def _dispatch(self, env: simpy.Environment, pe_txn: PeInternalTxn) -> Generator:
        from kernbench.common.pe_commands import CompositeCmd, PeCpuOverheadCmd

        yield from self.run(env, 0)  # scheduler overhead

        cmd = pe_txn.command

        # Simple command dispatch
        engine_suffix = self._CMD_DISPATCH.get(type(cmd))
        if engine_suffix is not None:
            yield self.out_ports[f"{self._pe_prefix}.{engine_suffix}"].put(pe_txn)
            return

        # CompositeCmd: generate plan and feed
        if isinstance(cmd, CompositeCmd):
            yield from self._dispatch_composite(env, pe_txn, cmd)
            return

        if isinstance(cmd, PeCpuOverheadCmd):
            yield env.timeout(cmd.cycles)
            pe_txn.done.succeed()
            return

        pe_txn.done.succeed()

    def _dispatch_composite(
        self, env: simpy.Environment, pe_txn: Any, cmd: Any,
    ) -> Generator:
        """Generate plan and enqueue to feeder. Non-blocking (ADR-0014 D6).

        Strict-FIFO RW gate (ADR-0065 D6.3): wait for overlapping in-flight
        composites, register this one's write set, retire on completion.
        Legacy composites (no rw_handles) pass through immediately.
        """
        from kernbench.components.builtin.pe_types import PipelineContext

        yield from self._hazard.admit(env, cmd, pe_txn.done)
        # Retire from the hazard tracker AND record the composite's numeric
        # op — both at completion (see _retire_on_done).
        env.process(self._retire_on_done(env, pe_txn.done, cmd))

        plan = self._generate_plan(cmd)

        self._pipeline_counter += 1
        # Prologue / post-loop single-shot stages each count as a tile for
        # completion (ADR-0065 D3). Empty for legacy composites → unchanged.
        n_pre = len(getattr(plan, "prologue_stages", ()))
        n_post = len(getattr(plan, "epilogue_stages", ()))
        ctx = PipelineContext(
            id=f"p{self._pipeline_counter}",
            total_tiles=len(plan.tiles) + n_pre + n_post,
            done_event=pe_txn.done,
        )

        # Enqueue to feeder — scheduler worker returns immediately
        assert self._pending_feeds is not None
        yield self._pending_feeds.put((plan, ctx))

    def _retire_on_done(
        self, env: simpy.Environment, done_event: Any, cmd: Any,
    ) -> Generator:
        """On composite completion: retire from the RW hazard tracker AND
        record its numeric op (ADR-0027 / ADR-0065 N1+N3).

        Recording at completion (not dispatch) gives the composite's GEMM a
        t_start AFTER its prologue MATH ops, so in data mode it reads their
        results (P, the rescaled O) instead of stale values. No-op without an
        op_logger → latency-only mode unchanged.
        """
        yield done_event
        self._hazard.retire(cmd.completion.id)
        self._on_process_start(env, cmd)
        self._on_process_end(env, cmd)

    def _feed_loop(self, env: simpy.Environment) -> Generator:
        """Single feeder process: FIFO command ordering (ADR-0014 D6).

        No tile feed interleaving between commands.
        Queue full → only this process blocks.
        """
        from kernbench.components.builtin.pe_types import TileToken

        assert self._pending_feeds is not None
        while True:
            plan, ctx = yield self._pending_feeds.get()
            tid = 0
            # Prologue single-shot stages first (ADR-0065 D3) — each a
            # standalone 1-stage tile that runs on its component and
            # completes (no inter-component routing).
            for st in getattr(plan, "prologue_stages", ()):
                yield from self._feed_single_stage(st, ctx, tid)
                tid += 1
            for tile in plan.tiles:
                first_stage = tile.stages[0]
                token = TileToken(
                    tile_id=tile.tile_id,
                    pipeline_ctx=ctx,
                    plan=tile,
                    stage_idx=0,
                    params=first_stage.params,
                )
                yield self.out_ports[first_stage.component].put(token)
            for st in getattr(plan, "epilogue_stages", ()):
                yield from self._feed_single_stage(st, ctx, tid)
                tid += 1

    def _feed_single_stage(self, stage: Any, ctx: Any, tile_id: int) -> Generator:
        """Feed one standalone stage as a 1-stage tile (prologue/post-loop)."""
        from kernbench.components.builtin.pe_types import TilePlan, TileToken

        token = TileToken(
            tile_id=tile_id, pipeline_ctx=ctx,
            plan=TilePlan(tile_id=tile_id, stages=(stage,)),
            stage_idx=0, params=stage.params,
        )
        yield self.out_ports[stage.component].put(token)

    def _generate_plan(self, cmd: Any) -> Any:
        """Generate a PipelinePlan from a flat-ops CompositeCmd (ADR-0065 D3).

        Position-scan for the GEMM (≤1) + prologue/epilogue placement lives
        in ``generate_plan_from_ops``; a legacy composite (GEMM at index 0,
        no prologue) produces the identical plan it did before.
        """
        from kernbench.components.builtin.tiling import generate_plan_from_ops

        return generate_plan_from_ops(
            ops=cmd.ops,
            tile_m=self.TILE_M, tile_k=self.TILE_K, tile_n=self.TILE_N,
            bytes_per_element=2,  # f16
            pe_prefix=self._pe_prefix,
        )
