"""
PGDAGRunner — The proof-gated DAG execution engine.

Orchestrates template execution: slot management, step sequencing,
proof collection, interruption handling, and slot closure.

Supports two DAG models:
  - Linear: steps run in declared order (simple, most templates)
  - Topological: steps run based on dependency graph (parallel branches)

The runner is product-agnostic.  Products provide:
  - DAGTemplate (step sequence or dependency graph)
  - StepFn implementations (the business logic)
  - Domain/frontier configuration

The kernel provides:
  - Slot open/close via SBN Gateway
  - Step sequencing with interruption checks
  - Proof collection + Merkle root computation
  - RunResult with full proof chain
"""

from __future__ import annotations

import logging
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Awaitable, Callable, Dict, List, Optional, Set

from .artifact import ProofArtifact, ProofStatus, merkle_root
from .envelope import StepEnvelopeBuilder
from .gate import IdempotencyGate
from .interrupt import InterruptionHandler
from .reducer import Reducer, ReducerStatus, RunState
from .runner import ProofWriter, StepContext, StepFn, StepRunner

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# WaitForEvent — pause DAG execution for external events
# ---------------------------------------------------------------------------

class WaitForEvent(Exception):
    """Raised by a step function to pause the DAG until an external event arrives.

    When a step raises this, the runner treats the node as *not failed* but
    *not completed* — the RunResult will carry ``waiting=True`` so the caller
    can persist the waiting state and an event bridge can resume the run later.

    Attributes:
        event_type: The event bus topic to wait for (e.g. "slot.completed").
        event_key:  An identifier the bridge uses to match the event
                    (e.g. the slot_id that must complete).
        node_id:    Populated automatically by the runner — the step that is waiting.
        partial_outputs: Any outputs the step wants merged before pausing.
    """

    def __init__(
        self,
        event_type: str,
        event_key: str,
        partial_outputs: Dict[str, Any] | None = None,
    ) -> None:
        self.event_type = event_type
        self.event_key = event_key
        self.node_id = ""  # set by the runner
        self.partial_outputs = partial_outputs or {}
        super().__init__(f"Waiting for {event_type} key={event_key}")


# ---------------------------------------------------------------------------
# DAG template definitions
# ---------------------------------------------------------------------------

@dataclass
class DAGNode:
    """A node in a DAG template.

    For linear templates, nodes are executed in list order.
    For topological templates, edges determine execution order.
    """
    id: str
    label: str = ""
    step_type: str = ""  # "ingestion" | "validation" | "computation" | "execution" | etc.
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class DAGEdge:
    """A directed edge in a DAG template.

    edge_type:
      - "sequence": source must complete before target starts
      - "fork": source fans out to multiple targets (parallel)
      - "join": multiple sources must complete before target starts

    The ``condition`` parameter is a backward-compatible alias for
    ``edge_type`` (used by Grid templates).
    """
    source: str
    target: str
    edge_type: str = "sequence"
    condition: str | None = None

    def __post_init__(self) -> None:
        if self.condition and self.edge_type == "sequence":
            self.edge_type = self.condition


@dataclass
class DAGTemplate:
    """A reusable workflow definition for PG-DAG execution.

    Products register templates with their step sequences.
    The PGDAGRunner executes templates by looking up registered StepFns.

    Two modes:
      - Linear: just nodes, executed in list order (no edges needed)
      - Topological: nodes + edges, executed by dependency resolution
    """
    name: str
    version: str = "1.0"
    nodes: List[DAGNode] = field(default_factory=list)
    edges: List[DAGEdge] = field(default_factory=list)
    domain: str = "general"
    metadata: Dict[str, Any] = field(default_factory=dict)

    @property
    def node_ids(self) -> List[str]:
        return [n.id for n in self.nodes]

    @property
    def is_linear(self) -> bool:
        """True if no edges defined — execute nodes in list order."""
        return len(self.edges) == 0

    def dependency_map(self) -> Dict[str, Set[str]]:
        """Build a map of node_id → set of prerequisite node_ids."""
        deps: Dict[str, Set[str]] = {n.id: set() for n in self.nodes}
        for edge in self.edges:
            if edge.target in deps:
                deps[edge.target].add(edge.source)
        return deps

    def ready_steps(self, completed: Set[str]) -> List[str]:
        """Return steps whose dependencies are all satisfied."""
        if self.is_linear:
            # Linear: return next uncompleted step
            for nid in self.node_ids:
                if nid not in completed:
                    return [nid]
            return []
        # Topological: return all steps with satisfied dependencies
        deps = self.dependency_map()
        ready = []
        for nid in self.node_ids:
            if nid not in completed and deps[nid].issubset(completed):
                ready.append(nid)
        return ready


# ---------------------------------------------------------------------------
# Node result (Grid-compatible per-step result shape)
# ---------------------------------------------------------------------------

@dataclass
class NodeResult:
    """Execution result for a single DAG node.

    Grid compatibility layer — maps proof artifacts to the simple
    success/outputs/error shape that Grid's _persist_dag_result expects.
    """
    node_id: str
    success: bool = False
    outputs: Dict[str, Any] = field(default_factory=dict)
    error: Optional[str] = None
    duration_ms: float = 0.0


# ---------------------------------------------------------------------------
# Run result
# ---------------------------------------------------------------------------

@dataclass
class RunResult:
    """Complete result of a DAG template execution."""
    run_id: str
    template_name: str
    status: ReducerStatus
    proofs: Dict[str, ProofArtifact]
    merkle_root_hash: str
    slot_id: Optional[str] = None
    slot_receipt_id: Optional[str] = None
    started_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    completed_at: Optional[datetime] = None
    # WaitForEvent support
    waiting: bool = False
    wait_event: Optional[WaitForEvent] = None
    final_outputs: Dict[str, Any] = field(default_factory=dict)
    duration_ms: float = 0.0

    @property
    def success(self) -> bool:
        """True if all expected steps completed successfully."""
        return self.status == ReducerStatus.COMPLETE

    @property
    def node_results(self) -> Dict[str, NodeResult]:
        """Convert proof chain to Grid-compatible NodeResult dict."""
        results: Dict[str, NodeResult] = {}
        for step_id, proof in self.proofs.items():
            results[step_id] = NodeResult(
                node_id=step_id,
                success=proof.succeeded,
                outputs={},
                error=proof.metrics.get("error") if proof.metrics else None,
                duration_ms=proof.metrics.get("duration_ms", 0.0) if proof.metrics else 0.0,
            )
        return results

    def to_dict(self) -> Dict[str, Any]:
        return {
            "run_id": self.run_id,
            "template_name": self.template_name,
            "status": self.status.value,
            "proofs": {k: v.to_dict() for k, v in self.proofs.items()},
            "merkle_root_hash": self.merkle_root_hash,
            "slot_id": self.slot_id,
            "slot_receipt_id": self.slot_receipt_id,
            "started_at": self.started_at.isoformat(),
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "waiting": self.waiting,
            "final_outputs": self.final_outputs,
        }


# Grid backward-compatibility alias
DAGResult = RunResult


# ---------------------------------------------------------------------------
# PGDAGRunner
# ---------------------------------------------------------------------------

class PGDAGRunner:
    """Proof-Gated DAG execution engine.

    Product-agnostic.  Products register step functions and execute templates.

    Usage:
        runner = PGDAGRunner(
            snapchore=sbn_client.snapchore,
            gateway=sbn_client.gateway,
            blocks=sbn_client.blocks,
            domain="treasury.execution",
            spec_version="1.0",
        )

        runner.register_steps({
            "acquire_rates": my_acquire_rates_fn,
            "seal_risk_envelope": my_seal_risk_fn,
            ...
        })

        result = await runner.execute(yield_capture_template, inputs={...}, actor="sp-engine-yield")
    """

    def __init__(
        self,
        snapchore: Any = None,
        gateway: Any = None,
        blocks: Any = None,
        *,
        domain: str = "general",
        spec_version: str = "1.0",
    ) -> None:
        self._snapchore = snapchore
        self._gateway = gateway
        self._domain = domain
        self._spec_version = spec_version
        self._steps: Dict[str, StepFn] = {}
        self._edge_emitter: Any = None  # Optional EdgeEmitter for Lattice edges

        # Build internal components
        from .envelope import SnapChoreCapture
        self._envelope_builder = StepEnvelopeBuilder(
            snapchore=snapchore if isinstance(snapchore, SnapChoreCapture) else None,
        )

        # Choose seal backend
        if snapchore is not None:
            from .runner import SbnSealBackend
            backend = SbnSealBackend(sbn_client=type("_SBN", (), {
                "snapchore": snapchore,
                "blocks": blocks,
            })())
        else:
            from .runner import LocalSealBackend
            backend = LocalSealBackend()

        self._writer = ProofWriter(backend)
        self._gate = IdempotencyGate(snapchore=snapchore)
        self._interruption = InterruptionHandler()

    @property
    def interruption_handler(self) -> InterruptionHandler:
        return self._interruption

    @property
    def idempotency_gate(self) -> IdempotencyGate:
        return self._gate

    # ── Step registration ─────────────────────────────────────────────

    def register_step(self, step_id: str, step_fn: StepFn) -> None:
        """Register a step function by ID."""
        self._steps[step_id] = step_fn

    def register_steps(self, steps: Dict[str, StepFn]) -> None:
        """Register multiple step functions at once."""
        self._steps.update(steps)

    # ── Template execution ────────────────────────────────────────────

    async def execute(
        self,
        template: DAGTemplate,
        inputs: Dict[str, Any],
        *,
        actor: str = "",
        config: Optional[Dict[str, Any]] = None,
        run_id: Optional[str] = None,
    ) -> RunResult:
        """Execute a DAG template and return a RunResult.

        1. Open SBN slot
        2. Resolve step order (linear or topological)
        3. For each step: check interruption → run → collect proof
        4. Reduce proofs → compute Merkle root
        5. Close SBN slot with attestation
        """
        run_id = run_id or f"run-{uuid.uuid4().hex[:12]}"
        run_config = config or {}
        started_at = datetime.now(timezone.utc)
        domain = template.domain or self._domain

        # Reset interruption handler for this run
        self._interruption.reset()

        # Open slot
        slot_id = await self._open_slot(run_id, template.name, actor)

        # Build step runner for this run
        step_runner = StepRunner(
            envelope_builder=self._envelope_builder,
            proof_writer=self._writer,
            gate=self._gate,
            domain=domain,
        )

        # Execute steps
        proofs: Dict[str, ProofArtifact] = {}
        current_outputs = dict(inputs)
        completed: Set[str] = set()

        while True:
            # Check for interruption between steps
            if self._interruption.is_triggered:
                last = list(completed)[-1] if completed else ""
                int_proof = self._interruption.build_interruption_proof(
                    last_completed_step=last,
                    run_id=run_id,
                    spec_version=self._spec_version,
                )
                proofs["interrupted"] = int_proof
                break

            # Get ready steps
            ready = template.ready_steps(completed)
            if not ready:
                break

            # Execute ready steps (sequentially for now; parallel TODO)
            for step_id in ready:
                if step_id not in self._steps:
                    logger.error("No registered function for step: %s", step_id)
                    proofs[step_id] = ProofArtifact(
                        step_id=step_id,
                        spec_version=self._spec_version,
                        inputs_fingerprint="",
                        outputs_fingerprint="",
                        proof_hash="",
                        block_ref="",
                        status=ProofStatus.FAILED,
                        metrics={"error": f"Step function not registered: {step_id}"},
                    )
                    completed.add(step_id)
                    continue

                # Build step inputs: run inputs + outputs from prior steps
                step_inputs = self._build_step_inputs(
                    step_id, current_outputs, proofs, template,
                )

                context = StepContext(
                    inputs=step_inputs,
                    prior_proofs=dict(proofs),
                    config=run_config,
                    actor=actor,
                    run_id=run_id,
                    slot_id=slot_id,
                    node_id=step_id,
                )

                # Execute step — catch WaitForEvent before proof sealing
                try:
                    proof = await step_runner.run(
                        step_id=step_id,
                        spec_version=self._spec_version,
                        step_fn=self._steps[step_id],
                        context=context,
                    )
                except WaitForEvent as wait:
                    # Step needs an external event — pause the DAG
                    wait.node_id = step_id
                    current_outputs.update(wait.partial_outputs)
                    return RunResult(
                        run_id=run_id,
                        template_name=template.name,
                        status=ReducerStatus.WAITING,
                        proofs=proofs,
                        merkle_root_hash="",
                        slot_id=slot_id,
                        started_at=started_at,
                        waiting=True,
                        wait_event=wait,
                        final_outputs=current_outputs,
                        duration_ms=(datetime.now(timezone.utc) - started_at).total_seconds() * 1000,
                    )

                proofs[step_id] = proof
                completed.add(step_id)

                # Emit Lattice edge for successful proofs
                if self._edge_emitter and proof.succeeded:
                    try:
                        await self._edge_emitter.emit(
                            step_id, proof,
                            run_id=run_id, template_name=template.name,
                        )
                    except Exception:
                        logger.warning("Edge emission failed for %s", step_id)

                # Thread outputs forward
                if proof.succeeded and context.inputs:
                    current_outputs.update(context.inputs)

                # If step failed, stop the run
                if not proof.succeeded:
                    break

            # Check if any step in this batch failed
            if any(
                proofs.get(s) and not proofs[s].succeeded
                for s in ready
                if s in proofs
            ):
                break

        # Reduce
        reducer = Reducer(expected_steps=template.node_ids, run_id=run_id)
        for proof in proofs.values():
            reducer.fold(proof)
        state = reducer.state()

        # Close slot
        receipt_id = await self._close_slot(
            slot_id, run_id, state.merkle_root_hash, proofs,
        )

        completed_at = datetime.now(timezone.utc)
        return RunResult(
            run_id=run_id,
            template_name=template.name,
            status=state.status,
            proofs=proofs,
            merkle_root_hash=state.merkle_root_hash,
            slot_id=slot_id,
            slot_receipt_id=receipt_id,
            started_at=started_at,
            completed_at=completed_at,
            final_outputs=current_outputs,
            duration_ms=(completed_at - started_at).total_seconds() * 1000,
        )

    # ── Slot management ───────────────────────────────────────────────

    async def _open_slot(
        self,
        run_id: str,
        template_name: str,
        actor: str,
    ) -> Optional[str]:
        """Open an SBN execution slot for this run."""
        if self._gateway is None:
            return f"local-{uuid.uuid4().hex[:12]}"
        try:
            result = self._gateway.create_slot({
                "slot_type": "execution",
                "run_id": run_id,
                "template": template_name,
                "actor": actor,
            })
            slot_id = result.get("slot_id", result.get("id", ""))
            logger.info("Opened slot %s for run %s", slot_id, run_id)
            return slot_id
        except Exception as exc:
            logger.warning("Failed to open SBN slot: %s — using local", exc)
            return f"local-{uuid.uuid4().hex[:12]}"

    async def _close_slot(
        self,
        slot_id: Optional[str],
        run_id: str,
        root: str,
        proofs: Dict[str, ProofArtifact],
    ) -> Optional[str]:
        """Close the SBN slot with Merkle root and request attestation."""
        if self._gateway is None or not slot_id or slot_id.startswith("local-"):
            return None
        try:
            result = self._gateway.close_slot(slot_id, {
                "run_id": run_id,
                "merkle_root": root,
                "step_count": len(proofs),
                "status": "complete" if all(p.succeeded for p in proofs.values()) else "partial",
            })
            receipt_id = result.get("receipt_id", result.get("attestation_id", ""))
            logger.info("Closed slot %s — receipt %s", slot_id, receipt_id)
            return receipt_id
        except Exception as exc:
            logger.warning("Failed to close SBN slot %s: %s", slot_id, exc)
            return None

    # ── Input threading ───────────────────────────────────────────────

    def _build_step_inputs(
        self,
        step_id: str,
        run_inputs: Dict[str, Any],
        proofs: Dict[str, ProofArtifact],
        template: DAGTemplate,
    ) -> Dict[str, Any]:
        """Build inputs for a step by merging run inputs + prior proof metadata."""
        inputs = dict(run_inputs)

        # Add proof hashes from dependencies
        if not template.is_linear:
            deps = template.dependency_map().get(step_id, set())
            for dep_id in deps:
                if dep_id in proofs:
                    inputs[f"{dep_id}_proof_hash"] = proofs[dep_id].proof_hash
                    inputs[f"{dep_id}_outputs_fingerprint"] = proofs[dep_id].outputs_fingerprint
        else:
            # Linear: thread all prior proof hashes
            for pid, proof in proofs.items():
                if proof.succeeded:
                    inputs[f"{pid}_proof_hash"] = proof.proof_hash

        return inputs
