"""
GuardPipeline — Stage-based composition primitive for llm-guard-kit components.
================================================================================

Wires the three-lane architecture (score → verify → repair) into a single
composable object. Each stage is conditional: run_if determines whether the
stage executes based on results so far.

Usage
-----
    pipeline = (
        GuardPipeline()
        .add_stage("adversarial", AdversarialChainDetector.load_default())
        .add_stage("score", AgentGuard(api_key=...))
        .add_stage("budget", ConformalRiskBudget(0.20), run_if=lambda r: r.final_risk > 0.3)
        # SelfHealer validated on RETRIEVAL_FAILURE (+11.1pp) and EXCESSIVE_SEARCH (+20pp)
        # in exp_selfhealer_ab (2026-03-20). Overall healed pass@1 = 6.1%.
        # Effect is modest — validate on your domain before relying on recovery in production.
        .add_stage("repair", SelfHealer(), run_if=lambda r: r.final_risk > 0.5)
    )
    result = pipeline.run(question, steps, final_answer)
    print(result.action)  # "proceed" | "rewrite" | "escalate" | "abort"
Purpose:    Wires score → verify → repair stages into a single composable object with conditional run_if predicates. Supports adversarial detection, risk scoring, budget allocation, and repair in a unified fluent API.
Exports:    GuardPipeline, PipelineResult
Added:      v0.33.0
Status:     Stable
"""
from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Literal, Optional, Tuple


@dataclass
class PipelineResult:
    stage_results:    Dict[str, Any]
    final_risk:       float
    action:           Literal["proceed", "rewrite", "escalate", "abort"]
    triggered_stages: List[str]
    skipped_stages:   List[str]
    latency_ms:       float
    metadata:         Dict[str, Any] = field(default_factory=dict)


def _extract_risk(result: Any) -> Optional[float]:
    """Duck-typed risk extraction from any component result."""
    if result is None:
        return None
    if isinstance(result, (int, float)):
        return float(result)
    risk = getattr(result, "risk_score", None)
    if risk is not None:
        return float(risk)
    if isinstance(result, dict):
        risk = result.get("risk_score")
        if risk is not None:
            return float(risk)
    return None


def _dispatch(name: str, component: Any, question: str, steps: list, final_answer: str) -> Any:
    """Call the first matching method on a component."""
    if hasattr(component, "score_chain"):
        return component.score_chain(question, steps, final_answer)
    if hasattr(component, "score"):
        return component.score(question, steps, final_answer)
    if hasattr(component, "run"):
        return component.run(question, steps, final_answer)
    raise TypeError(
        f"Component '{name}' has none of: score_chain, score, run. "
        f"Implement one of these methods to use it in a GuardPipeline."
    )


def _derive_action(
    risk: float,
    abort_threshold: float,
    escalate_threshold: float,
    rewrite_threshold: float,
) -> Literal["proceed", "rewrite", "escalate", "abort"]:
    if risk >= abort_threshold:
        return "abort"
    if risk >= escalate_threshold:
        return "escalate"
    if risk >= rewrite_threshold:
        return "rewrite"
    return "proceed"


class GuardPipeline:
    """Stage-based composition primitive. Stages execute in insertion order."""

    def __init__(
        self,
        abort_threshold:    float = 0.70,
        escalate_threshold: float = 0.50,
        rewrite_threshold:  float = 0.30,
    ) -> None:
        self._abort_threshold    = abort_threshold
        self._escalate_threshold = escalate_threshold
        self._rewrite_threshold  = rewrite_threshold
        # list of (name, component, run_if) tuples — NOT a dict
        self._stages: List[Tuple[str, Any, Optional[Callable]]] = []

    def add_stage(
        self,
        name: str,
        component: Any,
        run_if: Optional[Callable[[PipelineResult], bool]] = None,
    ) -> "GuardPipeline":
        """Add a stage. Returns self for fluent chaining."""
        existing = [s[0] for s in self._stages]
        if name in existing:
            raise ValueError(f"Stage name '{name}' already exists in this pipeline.")
        self._stages.append((name, component, run_if))
        return self

    def run(self, question: str, steps: list, final_answer: str) -> PipelineResult:
        """Execute all stages in insertion order, respecting run_if conditions."""
        t0 = time.perf_counter()

        current = PipelineResult(
            stage_results={},
            final_risk=0.0,
            action="proceed",
            triggered_stages=[],
            skipped_stages=[],
            latency_ms=0.0,
        )

        for name, component, run_if in self._stages:
            # Evaluate condition
            if run_if is not None and not run_if(current):
                current.skipped_stages.append(name)
                continue

            # Dispatch to component
            result = _dispatch(name, component, question, steps, final_answer)
            current.stage_results[name] = result
            current.triggered_stages.append(name)

            # Update final_risk
            risk = _extract_risk(result)
            if risk is not None:
                current.final_risk = risk

        # Derive action from final_risk
        current.action = _derive_action(
            current.final_risk,
            self._abort_threshold,
            self._escalate_threshold,
            self._rewrite_threshold,
        )
        current.latency_ms = (time.perf_counter() - t0) * 1000.0
        return current
