"""Tests for GuardPipeline — stage-based composition primitive."""
import sys
import pytest
import warnings
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from llm_guard.pipeline.guard_pipeline import GuardPipeline, PipelineResult


# ── Mock components ───────────────────────────────────────────────────────────

class MockScorer:
    """Has score_chain() — highest priority dispatch."""
    def __init__(self, risk):
        self.risk = risk
    def score_chain(self, question, steps, final_answer):
        from dataclasses import dataclass
        @dataclass
        class R:
            risk_score: float
        return R(self.risk)


class MockScoreOnly:
    """Has score() only — second priority."""
    def __init__(self, risk):
        self.risk = risk
    def score(self, question, steps, final_answer):
        return self.risk  # bare float


class MockRunOnly:
    """Has run() only — third priority."""
    def __init__(self, risk):
        self.risk = risk
    def run(self, question, steps, final_answer):
        return {"risk_score": self.risk}


class MockNoInterface:
    """Has none of the required methods."""
    pass


Q, STEPS, FA = "What year?", [{"thought": "check"}], "1912"


# ── Basic construction ────────────────────────────────────────────────────────

def test_add_stage_returns_self():
    p = GuardPipeline()
    result = p.add_stage("score", MockScorer(0.5))
    assert result is p


def test_duplicate_name_raises():
    p = GuardPipeline()
    p.add_stage("score", MockScorer(0.5))
    with pytest.raises(ValueError, match="already exists"):
        p.add_stage("score", MockScorer(0.3))


def test_no_interface_raises_type_error():
    p = GuardPipeline()
    p.add_stage("bad", MockNoInterface())
    with pytest.raises(TypeError, match="score_chain, score, run"):
        p.run(Q, STEPS, FA)


# ── Dispatch priority ─────────────────────────────────────────────────────────

def test_score_chain_takes_priority():
    p = GuardPipeline().add_stage("s", MockScorer(0.8))
    r = p.run(Q, STEPS, FA)
    assert abs(r.final_risk - 0.8) < 0.01


def test_score_method_used_when_no_score_chain():
    p = GuardPipeline().add_stage("s", MockScoreOnly(0.6))
    r = p.run(Q, STEPS, FA)
    assert abs(r.final_risk - 0.6) < 0.01


def test_run_method_used_as_fallback():
    p = GuardPipeline().add_stage("s", MockRunOnly(0.4))
    r = p.run(Q, STEPS, FA)
    assert abs(r.final_risk - 0.4) < 0.01


# ── Action derivation ─────────────────────────────────────────────────────────

def test_action_abort_above_threshold():
    p = GuardPipeline(abort_threshold=0.70)
    p.add_stage("s", MockScorer(0.85))
    r = p.run(Q, STEPS, FA)
    assert r.action == "abort"


def test_action_escalate():
    p = GuardPipeline(abort_threshold=0.70, escalate_threshold=0.50)
    p.add_stage("s", MockScorer(0.60))
    r = p.run(Q, STEPS, FA)
    assert r.action == "escalate"


def test_action_rewrite():
    p = GuardPipeline(abort_threshold=0.70, escalate_threshold=0.50, rewrite_threshold=0.30)
    p.add_stage("s", MockScorer(0.40))
    r = p.run(Q, STEPS, FA)
    assert r.action == "rewrite"


def test_action_proceed():
    p = GuardPipeline(rewrite_threshold=0.30)
    p.add_stage("s", MockScorer(0.10))
    r = p.run(Q, STEPS, FA)
    assert r.action == "proceed"


# ── run_if conditional execution ──────────────────────────────────────────────

def test_stage_skipped_when_run_if_false():
    p = GuardPipeline()
    p.add_stage("first", MockScorer(0.2))
    p.add_stage("second", MockScorer(0.9), run_if=lambda r: r.final_risk > 0.5)
    r = p.run(Q, STEPS, FA)
    assert "second" in r.skipped_stages
    assert abs(r.final_risk - 0.2) < 0.01


def test_stage_runs_when_run_if_true():
    p = GuardPipeline()
    p.add_stage("first", MockScorer(0.8))
    p.add_stage("second", MockScorer(0.9), run_if=lambda r: r.final_risk > 0.5)
    r = p.run(Q, STEPS, FA)
    assert "second" in r.triggered_stages


def test_first_stage_run_if_receives_default_result():
    """First stage's run_if receives a default PipelineResult with final_risk=0.0."""
    seen = {}
    def capture_run_if(result):
        seen["final_risk"] = result.final_risk
        return True
    p = GuardPipeline()
    p.add_stage("first", MockScorer(0.5), run_if=capture_run_if)
    p.run(Q, STEPS, FA)
    assert seen["final_risk"] == 0.0


# ── Result structure ──────────────────────────────────────────────────────────

def test_result_has_correct_fields():
    p = GuardPipeline().add_stage("s", MockScorer(0.5))
    r = p.run(Q, STEPS, FA)
    assert isinstance(r, PipelineResult)
    assert isinstance(r.stage_results, dict)
    assert isinstance(r.triggered_stages, list)
    assert isinstance(r.skipped_stages, list)
    assert r.latency_ms >= 0


def test_no_stages_returns_proceed():
    p = GuardPipeline()
    r = p.run(Q, STEPS, FA)
    assert r.final_risk == 0.0
    assert r.action == "proceed"
