"""PE_MATH: element-wise / reduction computation engine (ADR-0014 D1).

Handles both legacy PeInternalTxn (MathCmd) and pipeline TileToken.
In pipeline mode, receives token after fetch stage, computes SIMD, chains to next.

SIMD latency model (from pe_accel):
  cycles = ceil(num_elements / vector_width)
  latency_ns = cycles / clock_freq_ghz
"""
from __future__ import annotations

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

import simpy

from kernbench.components.base import PeEngineBase

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


class PeMathComponent(PeEngineBase):
    """PE_MATH: SIMD/Vector unit (ADR-0014 D1).

    In pipeline mode: pure compute — register data already fetched.
    In legacy mode: handles PeInternalTxn(MathCmd) with shared accel_slot.
    """

    def __init__(self, node: Node, ctx: ComponentContext | None = None) -> None:
        super().__init__(node, ctx)
        self._accel: simpy.Resource | None = None
        self._vector_width: int = int(node.attrs.get("vector_width", 256))
        self._clock_freq: float = float(node.attrs.get("clock_freq_ghz", 1.0))

    def init_resources(self, env: simpy.Environment) -> None:
        resource_name = self.node.attrs.get("shared_resource")
        if resource_name and self.ctx:
            self._accel = self.ctx.get_shared_resource(
                env, f"{self._pe_prefix}.{resource_name}"
            )

    def _compute_ns(self, num_elements: int) -> float:
        """SIMD latency (pe_accel model)."""
        if self._vector_width > 0 and self._clock_freq > 0 and num_elements > 0:
            cycles = ceil(num_elements / self._vector_width)
            return cycles / self._clock_freq
        return float(self.node.attrs.get("overhead_ns", 0.0))

    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
        from kernbench.components.builtin.pe_types import TileToken

        while True:
            msg: Any = yield self._inbox.get()
            if isinstance(msg, TileToken):
                env.process(self._pipeline_process(env, msg))
            elif isinstance(msg, PeInternalTxn):
                env.process(self._handle_with_hooks(env, msg))
            else:
                env.process(self._forward_txn(env, msg))

    def _pipeline_process(self, env: simpy.Environment, token: Any) -> Generator:
        """Pipeline mode: pure SIMD compute, then self-route.

        Handles consecutive same-component MATH stages (e.g. chained
        epilogue ops bias→relu) by looping locally before forwarding,
        mirroring the pattern in PeDmaComponent._pipeline_process.
        """
        yield from self._do_compute(env, token)

        next_stage = token.advance()
        while next_stage is not None and next_stage.component == self.node.id:
            yield from self._do_compute(env, token)
            next_stage = token.advance()

        if next_stage is not None:
            yield self.out_ports[next_stage.component].put(token)
        else:
            token.pipeline_ctx.complete_tile()

    def _do_compute(self, env: simpy.Environment, token: Any) -> Generator:
        """Single MATH stage: latency model + op_log bracketing."""
        self._on_process_start(env, token)
        num_elements = token.params.get("num_elements", 0)
        if self._accel:
            with self._accel.request() as req:
                yield req
                yield env.timeout(self._compute_ns(num_elements))
        else:
            yield env.timeout(self._compute_ns(num_elements))
        self._on_process_end(env, token)

    def handle_command(self, env: simpy.Environment, pe_txn: PeInternalTxn) -> Generator:
        """PeInternalTxn handling for standalone MathCmd / CopyCmd.

        Latency = max(overhead_ns, _compute_ns(num_elements)):
        - overhead_ns: fixed per-invocation setup cost (from node attrs).
        - _compute_ns: SIMD cycle-based model (from vector_width + clock_freq).
        The larger of the two dominates (setup-bound vs compute-bound).

        CopyCmd (ADR-0063 §D3.1): vector-engine on-chip byte copy; cost
        model = _compute_ns(prod(dst.shape)).
        """
        from kernbench.common.pe_commands import CopyCmd, MathCmd
        import math as _math

        cmd = pe_txn.command
        num_elements = 0
        if isinstance(cmd, MathCmd) and cmd.out.shape:
            num_elements = _math.prod(cmd.out.shape)
        elif isinstance(cmd, CopyCmd) and cmd.dst.shape:
            num_elements = _math.prod(cmd.dst.shape)

        overhead_ns = float(self.node.attrs.get("overhead_ns", 0.0))
        compute_ns = self._compute_ns(num_elements)
        ns = max(overhead_ns, compute_ns)

        if self._accel:
            with self._accel.request() as req:
                yield req
                yield env.timeout(ns)
        else:
            yield env.timeout(ns)
        pe_txn.done.succeed()

    def _forward_txn(self, env: simpy.Environment, txn: Any) -> Generator:
        if self._accel:
            with self._accel.request() as req:
                yield req
                yield from super()._forward_txn(env, txn)
        else:
            yield from super()._forward_txn(env, txn)
