"""
Reducer — Derive run state by folding over proofs, not mutable flags.

The Reducer scans all ProofArtifacts from a run and produces a RunState:
  COMPLETE:     all expected steps sealed → Merkle root computed
  PARTIAL:      steps 1..N sealed, N+1 missing → "pending from N+1"
  FAILED:       sealed steps + failure proof → failure evidence
  INTERRUPTED:  sealed steps + interruption proof → NUMA signal ref

Incremental fold: call fold() per proof as they arrive, or fold_many()
with a batch.  state() returns the current reduction at any time.

The DB is a read cache.  If it drifts, replay proofs.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Dict, List, Optional, Set

from .artifact import ProofArtifact, ProofStatus, merkle_root


# ---------------------------------------------------------------------------
# Reducer status
# ---------------------------------------------------------------------------

class ReducerStatus(str, Enum):
    COMPLETE = "complete"
    PARTIAL = "partial"
    FAILED = "failed"
    INTERRUPTED = "interrupted"


# ---------------------------------------------------------------------------
# Run state
# ---------------------------------------------------------------------------

@dataclass
class RunState:
    """Reduced state of a single DAG run, derived entirely from proofs."""
    run_id: str
    status: ReducerStatus
    completed_steps: List[str] = field(default_factory=list)
    failed_steps: List[str] = field(default_factory=list)
    interrupted_steps: List[str] = field(default_factory=list)
    pending_from: Optional[str] = None
    last_completed_step: Optional[str] = None
    proofs: Dict[str, ProofArtifact] = field(default_factory=dict)
    merkle_root_hash: str = ""
    metrics: Dict[str, Any] = field(default_factory=dict)
    reduced_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))

    @property
    def all_sealed_steps(self) -> Set[str]:
        """All steps that have any proof (completed, failed, or interrupted)."""
        return set(self.completed_steps) | set(self.failed_steps) | set(self.interrupted_steps)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "run_id": self.run_id,
            "status": self.status.value,
            "completed_steps": self.completed_steps,
            "failed_steps": self.failed_steps,
            "interrupted_steps": self.interrupted_steps,
            "pending_from": self.pending_from,
            "last_completed_step": self.last_completed_step,
            "merkle_root_hash": self.merkle_root_hash,
            "metrics": self.metrics,
            "reduced_at": self.reduced_at.isoformat(),
        }


# ---------------------------------------------------------------------------
# Reducer
# ---------------------------------------------------------------------------

class Reducer:
    """Incremental proof-folding reducer.

    Usage:
        reducer = Reducer(expected_steps=["acquire_rates", "seal_risk", ...], run_id="run-123")

        for proof in proofs:
            reducer.fold(proof)

        state = reducer.state()
        # state.status == ReducerStatus.COMPLETE if all steps sealed
    """

    def __init__(
        self,
        expected_steps: List[str],
        run_id: str = "",
    ) -> None:
        self._expected = list(expected_steps)
        self._run_id = run_id
        self._completed: List[str] = []
        self._failed: List[str] = []
        self._interrupted: List[str] = []
        self._proofs: Dict[str, ProofArtifact] = {}
        self._metrics: Dict[str, Any] = {}

    def fold(self, proof: ProofArtifact) -> RunState:
        """Fold a single proof into the current state."""
        self._proofs[proof.step_id] = proof

        if proof.status == ProofStatus.COMPLETED:
            if proof.step_id not in self._completed:
                self._completed.append(proof.step_id)
        elif proof.status == ProofStatus.FAILED:
            if proof.step_id not in self._failed:
                self._failed.append(proof.step_id)
        elif proof.status == ProofStatus.INTERRUPTED:
            if proof.step_id not in self._interrupted:
                self._interrupted.append(proof.step_id)

        # Merge metrics
        if proof.metrics:
            self._metrics.update(proof.metrics)

        return self.state()

    def fold_many(self, proofs: List[ProofArtifact]) -> RunState:
        """Fold multiple proofs at once."""
        for proof in proofs:
            self.fold(proof)
        return self.state()

    def state(self) -> RunState:
        """Compute current RunState from folded proofs."""
        sealed = set(self._completed) | set(self._failed) | set(self._interrupted)

        # Determine status
        if self._interrupted:
            status = ReducerStatus.INTERRUPTED
        elif self._failed:
            status = ReducerStatus.FAILED
        elif sealed >= set(self._expected):
            status = ReducerStatus.COMPLETE
        else:
            status = ReducerStatus.PARTIAL

        # Find first pending step
        pending_from = None
        if status == ReducerStatus.PARTIAL:
            for step in self._expected:
                if step not in sealed:
                    pending_from = step
                    break

        # Merkle root over completed proof hashes (in expected order)
        proof_hashes = [
            self._proofs[s].proof_hash
            for s in self._expected
            if s in self._proofs
        ]
        root = merkle_root(proof_hashes) if proof_hashes else ""

        return RunState(
            run_id=self._run_id,
            status=status,
            completed_steps=list(self._completed),
            failed_steps=list(self._failed),
            interrupted_steps=list(self._interrupted),
            pending_from=pending_from,
            last_completed_step=self._completed[-1] if self._completed else None,
            proofs=dict(self._proofs),
            merkle_root_hash=root,
            metrics=dict(self._metrics),
        )

    def reset(self) -> None:
        """Reset the reducer for a new run."""
        self._completed.clear()
        self._failed.clear()
        self._interrupted.clear()
        self._proofs.clear()
        self._metrics.clear()
