from __future__ import annotations
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import numpy as np
import pytest


# ── helpers ──────────────────────────────────────────────────────────────────

def _make_run(n_search: int, label: int) -> dict:
    """label: 1=wrong, 0=correct"""
    steps = [
        {"thought": "need info", "action_type": "Search",
         "action_arg": "query", "observation": "some text found"}
        for _ in range(n_search)
    ]
    steps.append({"thought": "done", "action_type": "Finish",
                  "action_arg": "answer", "observation": ""})
    return {"question": "Who did X?", "steps": steps,
            "final_answer": "answer", "label": label}


def _make_runs(n=10):
    """5 correct (label=0), 5 wrong (label=1), varying step counts."""
    return (
        [_make_run(1, 0), _make_run(2, 0), _make_run(1, 0),
         _make_run(3, 0), _make_run(2, 0)] +
        [_make_run(4, 1), _make_run(5, 1), _make_run(3, 1),
         _make_run(6, 1), _make_run(4, 1)]
    )


# ═══════════════════════════════════════════════════════════════════════════
# SECTION 1 — PartialChainCalibrator
# ═══════════════════════════════════════════════════════════════════════════

class TestPartialChainCalibrator:

    def test_fit_creates_calibrators_for_depths(self):
        from llm_guard.agent_guard import PartialChainCalibrator, AgentGuard
        guard = AgentGuard()
        cal = PartialChainCalibrator(guard)
        cal.fit_partial(_make_runs())
        assert hasattr(cal, "calibrators_")
        assert set(cal.calibrators_.keys()) >= {1, 2, 3}

    def test_calibrate_returns_float_in_range(self):
        from llm_guard.agent_guard import PartialChainCalibrator, AgentGuard
        guard = AgentGuard()
        cal = PartialChainCalibrator(guard)
        cal.fit_partial(_make_runs())
        score = cal.calibrate(0.6, depth=2)
        assert 0.0 <= score <= 1.0

    def test_calibrate_fallback_for_unknown_depth(self):
        from llm_guard.agent_guard import PartialChainCalibrator, AgentGuard
        guard = AgentGuard()
        cal = PartialChainCalibrator(guard)
        cal.fit_partial(_make_runs())
        raw = 0.42
        result = cal.calibrate(raw, depth=9)
        assert result == raw  # falls back to raw score

    def test_agent_guard_fit_partial_calibration_stores_attr(self):
        from llm_guard.agent_guard import AgentGuard
        guard = AgentGuard()
        guard.fit_partial_calibration(_make_runs())
        assert hasattr(guard, "_partial_cal")
        assert guard._partial_cal is not None

    def test_stream_guard_uses_partial_cal_when_flag_set(self):
        from llm_guard.agent_guard import AgentGuard
        guard = AgentGuard()
        guard.fit_partial_calibration(_make_runs())
        steps = _make_run(2, 0)["steps"][:2]
        result = guard.stream_guard(
            "Who did X?", steps,
            use_partial_calibration=True,
        )
        # Just check it returns without error and has expected fields
        assert hasattr(result, "abort")
        assert hasattr(result, "risk_at_step")

    def test_calibrate_depth_greater_than_3_returns_raw(self):
        from llm_guard.agent_guard import PartialChainCalibrator, AgentGuard
        guard = AgentGuard()
        cal = PartialChainCalibrator(guard)
        cal.fit_partial(_make_runs())
        # depth > 3: raw score returned even though depth-3 calibrator exists
        raw = 0.77
        result = cal.calibrate(raw, depth=5)
        assert result == raw


# ═══════════════════════════════════════════════════════════════════════════
# SECTION 2 — langchain_trace_to_steps
# ═══════════════════════════════════════════════════════════════════════════

class TestLangchainTraceToSteps:

    def _make_action(self, tool="Search", tool_input="query", log="Thought: need info\nAction: Search\nAction Input: query"):
        """Create a mock AgentAction."""
        from unittest.mock import MagicMock
        action = MagicMock()
        action.tool = tool
        action.tool_input = tool_input
        action.log = log
        return action

    def test_empty_intermediate_steps_returns_empty_list(self):
        from llm_guard.integrations.langchain import langchain_trace_to_steps
        result = langchain_trace_to_steps([])
        assert result == []

    def test_single_step_produces_correct_step_dict(self):
        from llm_guard.integrations.langchain import langchain_trace_to_steps
        action = self._make_action(tool="Search", tool_input="Einstein birthday",
                                   log="Thought: look it up\nAction: Search")
        steps = langchain_trace_to_steps([(action, "Einstein born 1879")])
        assert len(steps) == 1
        s = steps[0]
        assert s["action_type"] == "Search"
        assert s["action_arg"] == "Einstein birthday"
        assert s["observation"] == "Einstein born 1879"
        assert "thought" in s

    def test_thought_extracted_from_log(self):
        from llm_guard.integrations.langchain import langchain_trace_to_steps
        action = self._make_action(log="Thought: this is my thought\nAction: Search")
        steps = langchain_trace_to_steps([(action, "obs")])
        assert steps[0]["thought"] == "this is my thought"

    def test_final_answer_appended_as_finish_step(self):
        from llm_guard.integrations.langchain import langchain_trace_to_steps
        action = self._make_action()
        steps = langchain_trace_to_steps([(action, "obs")], final_answer="42")
        assert steps[-1]["action_type"] == "Finish"
        assert steps[-1]["action_arg"] == "42"
        assert steps[-1]["observation"] == ""

    def test_dict_tool_input_converted_to_string(self):
        from llm_guard.integrations.langchain import langchain_trace_to_steps
        action = self._make_action(tool_input={"query": "birth year"})
        steps = langchain_trace_to_steps([(action, "obs")])
        assert isinstance(steps[0]["action_arg"], str)
        assert "birth year" in steps[0]["action_arg"]


# ═══════════════════════════════════════════════════════════════════════════
# SECTION 3 — DirectGenerationGuard
# ═══════════════════════════════════════════════════════════════════════════

class TestDirectGenerationGuard:

    def test_score_returns_direct_gen_result(self):
        from llm_guard.direct_generation_guard import DirectGenerationGuard
        guard = DirectGenerationGuard()
        result = guard.score("What is 2+2?", "4.", use_ptrue=False)
        assert hasattr(result, "risk_score")
        assert hasattr(result, "confidence_tier")
        assert hasattr(result, "lexical_risk")
        assert hasattr(result, "signals")

    def test_risk_score_in_range(self):
        from llm_guard.direct_generation_guard import DirectGenerationGuard
        guard = DirectGenerationGuard()
        result = guard.score("What?", "I'm not sure, maybe, possibly.", use_ptrue=False)
        assert 0.0 <= result.risk_score <= 1.0

    def test_hedge_heavy_answer_has_higher_risk(self):
        from llm_guard.direct_generation_guard import DirectGenerationGuard
        guard = DirectGenerationGuard()
        confident = guard.score("Who was Einstein?", "Albert Einstein was a physicist born in 1879.", use_ptrue=False)
        hedging = guard.score("Who was Einstein?", "I'm not sure but possibly he was maybe a scientist I think.", use_ptrue=False)
        assert hedging.risk_score > confident.risk_score

    def test_confidence_tier_matches_risk_score(self):
        from llm_guard.direct_generation_guard import DirectGenerationGuard, DirectGenResult
        guard = DirectGenerationGuard()
        r_high = guard.score("Q?", "Clear and complete answer with full stop.", use_ptrue=False)
        # LOW tier when risk >= 0.70
        # Manually test tier logic
        from llm_guard.direct_generation_guard import _risk_to_tier
        assert _risk_to_tier(0.3) == "HIGH"
        assert _risk_to_tier(0.6) == "MEDIUM"
        assert _risk_to_tier(0.8) == "LOW"

    def test_zero_cost_mode_ptrue_score_is_nan(self):
        from llm_guard.direct_generation_guard import DirectGenerationGuard
        import math
        guard = DirectGenerationGuard()
        result = guard.score("Q?", "answer.", use_ptrue=False)
        assert math.isnan(result.ptrue_score)


# ═══════════════════════════════════════════════════════════════════════════
# SECTION 4 — FunctionCallGuard
# ═══════════════════════════════════════════════════════════════════════════

class TestFunctionCallGuard:

    def _openai_messages(self):
        return [
            {"role": "user", "content": "Who was Einstein?"},
            {"role": "assistant", "content": "Let me search.",
             "tool_calls": [{"function": {"name": "search",
                                          "arguments": '{"query": "Einstein biography"}'}}]},
            {"role": "tool", "content": "Einstein was a physicist born 1879."},
            {"role": "assistant", "content": None,
             "tool_calls": [{"function": {"name": "finish",
                                          "arguments": '{"answer": "German physicist"}'}}]},
            {"role": "tool", "content": "German physicist"},
        ]

    def _claude_messages(self):
        return [
            {"role": "user", "content": "Who was Einstein?"},
            {"role": "assistant", "content": [
                {"type": "text", "text": "Let me search."},
                {"type": "tool_use", "name": "search",
                 "input": {"query": "Einstein biography"}},
            ]},
            {"role": "user", "content": [
                {"type": "tool_result", "content": "Einstein was a physicist born 1879."},
            ]},
        ]

    def test_parse_tool_calls_openai_format(self):
        from llm_guard.function_call_guard import FunctionCallGuard
        guard = FunctionCallGuard()
        steps = guard._parse_tool_calls(self._openai_messages())
        assert len(steps) >= 1
        assert "action_type" in steps[0]
        assert steps[0]["action_type"] == "search"
        assert "observation" in steps[0]

    def test_parse_tool_calls_claude_format(self):
        from llm_guard.function_call_guard import FunctionCallGuard
        guard = FunctionCallGuard()
        steps = guard._parse_tool_calls(self._claude_messages())
        assert len(steps) >= 1
        assert steps[0]["action_type"] == "search"
        assert "Einstein" in steps[0]["observation"]

    def test_score_returns_chain_trust_result(self):
        from llm_guard.function_call_guard import FunctionCallGuard
        from llm_guard.agent_guard import ChainTrustResult
        guard = FunctionCallGuard()
        result = guard.score(self._openai_messages(), question="Who was Einstein?")
        assert isinstance(result, ChainTrustResult)
        assert 0.0 <= result.risk_score <= 1.0
        assert result.confidence_tier in ("HIGH", "MEDIUM", "LOW")

    def test_score_empty_messages_returns_result(self):
        from llm_guard.function_call_guard import FunctionCallGuard
        guard = FunctionCallGuard()
        result = guard.score([], question="Q?")
        assert hasattr(result, "risk_score")

    def test_parse_extracts_arguments_string(self):
        from llm_guard.function_call_guard import FunctionCallGuard
        guard = FunctionCallGuard()
        steps = guard._parse_tool_calls(self._openai_messages())
        assert isinstance(steps[0]["action_arg"], str)

    def test_fit_raises_on_too_few_chains(self):
        from llm_guard.function_call_guard import FunctionCallGuard
        guard = FunctionCallGuard()
        with pytest.raises(ValueError, match="20"):
            guard.fit([{"messages": self._openai_messages(), "label": 0}] * 5)

    def test_fit_sets_tce_attribute(self):
        from llm_guard.function_call_guard import FunctionCallGuard
        guard = FunctionCallGuard()
        # Build 20 synthetic labeled chains: 10 correct, 10 wrong
        chains = (
            [{"messages": self._openai_messages(), "label": 0}] * 10 +
            [{"messages": self._claude_messages(), "label": 1}] * 10
        )
        guard.fit(chains)
        assert guard._tce is not None

    def test_score_after_fit_blends_tce_risk(self):
        from llm_guard.function_call_guard import FunctionCallGuard
        guard = FunctionCallGuard()
        chains = (
            [{"messages": self._openai_messages(), "label": 0}] * 10 +
            [{"messages": self._claude_messages(), "label": 1}] * 10
        )
        guard.fit(chains)
        result = guard.score(self._openai_messages(), "What is the capital of France?")
        assert "tce_risk" in result.behavioral_components
        assert 0.0 <= result.risk_score <= 1.0


# ═══════════════════════════════════════════════════════════════════════════
# SECTION 5 — MultiTurnGuard
# ═══════════════════════════════════════════════════════════════════════════

class TestMultiTurnGuard:

    def _conv(self, n_turns=3, quality="good"):
        """Build a simple conversation list."""
        msgs = []
        for i in range(n_turns):
            msgs.append({"role": "user", "content": f"Question {i}?"})
            if quality == "good":
                msgs.append({"role": "assistant",
                             "content": f"Clear complete answer to question {i}."})
            else:
                msgs.append({"role": "assistant",
                             "content": f"I'm not sure maybe possibly uncertain {i}."})
        return msgs

    def test_empty_conversation_returns_result(self):
        from llm_guard.multi_turn_guard import MultiTurnGuard
        guard = MultiTurnGuard()
        result = guard.score_turn([], use_ptrue=False)
        assert hasattr(result, "risk_score")
        assert result.window_size == 0

    def test_score_turn_returns_multi_turn_result(self):
        from llm_guard.multi_turn_guard import MultiTurnGuard, MultiTurnResult
        guard = MultiTurnGuard()
        result = guard.score_turn(self._conv(3), use_ptrue=False)
        assert isinstance(result, MultiTurnResult)
        assert 0.0 <= result.risk_score <= 1.0
        assert result.confidence_tier in ("HIGH", "MEDIUM", "LOW")

    def test_hedging_conversation_has_higher_risk(self):
        from llm_guard.multi_turn_guard import MultiTurnGuard
        guard = MultiTurnGuard()
        good = guard.score_turn(self._conv(3, "good"), use_ptrue=False)
        bad = guard.score_turn(self._conv(3, "bad"), use_ptrue=False)
        assert bad.risk_score > good.risk_score

    def test_drift_detected_when_quality_declining(self):
        from llm_guard.multi_turn_guard import MultiTurnGuard
        guard = MultiTurnGuard()
        # Manually inject a declining score sequence via the internal method
        result = guard._score_from_turn_scores([0.9, 0.7, 0.5, 0.3], window=4)
        assert result.drift_detected is True

    def test_window_limits_turns_scored(self):
        from llm_guard.multi_turn_guard import MultiTurnGuard
        guard = MultiTurnGuard()
        conv = self._conv(10, "good")
        result = guard.score_turn(conv, window=3, use_ptrue=False)
        assert result.window_size <= 3
