"""PE pipeline types for ADR-0014 D6: TileToken, TilePlan, Stage, PipelineContext.

These types are used by the PE_SCHEDULER and all PE engine components
for tile-based pipeline execution with self-routing.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum, auto
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    import simpy


# ── Stage types ──────────────────────────────────────────────────────


class StageType(Enum):
    DMA_READ = auto()
    FETCH = auto()
    GEMM = auto()
    MATH = auto()
    STORE = auto()
    DMA_WRITE = auto()


@dataclass
class Stage:
    """One stage in a tile's execution plan."""

    stage_type: StageType
    component: str      # topology node ID (e.g. "sip0.cube0.pe0.pe_dma")
    params: dict = field(default_factory=dict)


# ── Plan ─────────────────────────────────────────────────────────────


@dataclass
class TilePlan:
    """Execution plan for a single tile (immutable stage sequence)."""

    tile_id: int
    stages: tuple[Stage, ...]


@dataclass
class PipelinePlan:
    """Full pipeline plan for one CompositeCmd."""

    tiles: list[TilePlan]
    # Metadata for metrics
    m_tiles: int = 0
    k_tiles: int = 0
    n_tiles: int = 0
    # Flat-ops (ADR-0065 D3): single-shot MATH stages placed before / after
    # the GEMM tile loop (recipe prologue / KERNEL-scope post-loop ops).
    # Informational — execution folds these into the first / last tile so
    # the feeder + completion counting stay unchanged.
    prologue_stages: tuple[Stage, ...] = ()
    epilogue_stages: tuple[Stage, ...] = ()


# ── Pipeline Context ─────────────────────────────────────────────────


@dataclass
class PipelineContext:
    """Tracks completion of a pipeline (exactly-once contract).

    Each tile's last stage calls complete_tile() exactly once.
    When all tiles complete, done_event.succeed() is called.
    """

    id: str
    total_tiles: int
    completed_tiles: int = 0
    done_event: Any = None  # simpy.Event

    def complete_tile(self) -> None:
        self.completed_tiles += 1
        if self.completed_tiles == self.total_tiles:
            if self.done_event is not None:
                self.done_event.succeed()


# ── TileToken ────────────────────────────────────────────────────────


@dataclass
class TileToken:
    """Self-routing tile token passed between PE components (ADR-0014 D6).

    Single-owner: only one component holds this token at any time.
    params is a cache of plan.stages[stage_idx].params (canonical source).
    """

    tile_id: int
    pipeline_ctx: PipelineContext
    plan: TilePlan
    stage_idx: int
    params: dict = field(default_factory=dict)
    data_op: bool = True     # op_log recording target (ADR-0020)

    @property
    def current_stage(self) -> Stage:
        return self.plan.stages[self.stage_idx]

    @property
    def has_next_stage(self) -> bool:
        return self.stage_idx + 1 < len(self.plan.stages)

    def advance(self) -> Stage | None:
        """Advance to next stage. Returns next Stage or None if last."""
        self.stage_idx += 1
        if self.stage_idx < len(self.plan.stages):
            next_stage = self.plan.stages[self.stage_idx]
            self.params = next_stage.params
            return next_stage
        return None
