"""
Tests for v0.43.0 features:
  - MCP retrieval_audit tool (dispatch logic, not real MCP transport)
  - Latency benchmark utility functions (unit tests for helpers, not full benchmark)
"""

import json
import pytest
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))


# ── MCP retrieval_audit tool ──────────────────────────────────────────────────

class TestMCPRetrievalAuditTool:
    """Test the retrieval_audit dispatch logic in the MCP server."""

    def _dispatch_retrieval_audit(self, question: str, steps: list) -> dict:
        """Exercise the dispatch path for retrieval_audit without MCP transport."""
        from llm_guard import AgentGuard

        guard = AgentGuard()

        # Mirror the dispatch logic from server.py
        rq = guard.retrieval_quality(question, steps)

        search_steps = [s for s in steps if s.get("action_type") == "Search"
                        and s.get("observation", "").strip()]
        step_details = []
        for i, (sim, step) in enumerate(zip(rq.get("per_step", []), search_steps)):
            step_details.append({
                "step_index": i,
                "query":      step.get("action_arg", "")[:120],
                "cosine_sim": round(float(sim), 4),
                "relevant":   float(sim) > 0.25,
            })

        return {
            "quality_label":    rq["quality_label"],
            "mean_sim":         rq["mean_sim"],
            "max_sim":          rq["max_sim"],
            "min_sim":          rq["min_sim"],
            "var_sim":          rq["var_sim"],
            "n_relevant_steps": rq["n_relevant"],
            "n_search_steps":   len(search_steps),
            "per_step_details": step_details,
            "interpretation": (
                f"Retrieval quality is {rq['quality_label']}. "
                f"{rq['n_relevant']}/{len(search_steps)} steps returned relevant observations "
                f"(cosine_sim > 0.25). "
                + ("Consider rephrasing search queries to be more specific. "
                   if rq["mean_sim"] < 0.35 else
                   "At least some observations are topically relevant. "
                   if rq["mean_sim"] < 0.50 else
                   "Retrieval looks healthy — failure likely in reasoning, not search.")
            ),
        }

    QUESTION = "When was the Battle of Hastings?"
    STEPS = [
        {
            "thought": "Search for the battle date.",
            "action_type": "Search",
            "action_arg": "Battle of Hastings 1066",
            "observation": "The Battle of Hastings was fought on 14 October 1066.",
        },
        {
            "thought": "I have the answer.",
            "action_type": "Finish",
            "action_arg": "1066",
            "observation": "",
        },
    ]

    def test_returns_expected_keys(self):
        result = self._dispatch_retrieval_audit(self.QUESTION, self.STEPS)
        for key in [
            "quality_label", "mean_sim", "max_sim", "min_sim", "var_sim",
            "n_relevant_steps", "n_search_steps", "per_step_details", "interpretation"
        ]:
            assert key in result, f"Missing key: {key}"

    def test_quality_label_valid(self):
        result = self._dispatch_retrieval_audit(self.QUESTION, self.STEPS)
        assert result["quality_label"] in {"GOOD", "FAIR", "POOR"}

    def test_n_search_steps_correct(self):
        result = self._dispatch_retrieval_audit(self.QUESTION, self.STEPS)
        assert result["n_search_steps"] == 1  # only 1 Search step with observation

    def test_per_step_details_length(self):
        result = self._dispatch_retrieval_audit(self.QUESTION, self.STEPS)
        assert len(result["per_step_details"]) == 1

    def test_per_step_cosine_sim_in_range(self):
        result = self._dispatch_retrieval_audit(self.QUESTION, self.STEPS)
        for step in result["per_step_details"]:
            assert 0.0 <= step["cosine_sim"] <= 1.0
            assert isinstance(step["relevant"], bool)

    def test_interpretation_is_string(self):
        result = self._dispatch_retrieval_audit(self.QUESTION, self.STEPS)
        assert isinstance(result["interpretation"], str)
        assert len(result["interpretation"]) > 20

    def test_empty_steps_returns_poor(self):
        result = self._dispatch_retrieval_audit(self.QUESTION, [])
        assert result["quality_label"] == "POOR"
        assert result["n_search_steps"] == 0
        assert result["per_step_details"] == []

    def test_finish_only_steps_returns_poor(self):
        steps = [{"thought": "x", "action_type": "Finish",
                  "action_arg": "done", "observation": ""}]
        result = self._dispatch_retrieval_audit(self.QUESTION, steps)
        assert result["quality_label"] == "POOR"
        assert result["n_search_steps"] == 0

    def test_high_relevance_chain(self):
        """On-topic chain should score FAIR or better."""
        steps = [
            {
                "thought": "Let me search for this.",
                "action_type": "Search",
                "action_arg": "Battle of Hastings date year",
                "observation": "Battle of Hastings was in 1066, fought by William the Conqueror.",
            }
        ]
        result = self._dispatch_retrieval_audit(self.QUESTION, steps)
        assert result["mean_sim"] >= 0.25

    def test_off_topic_chain_low_sim(self):
        """Off-topic search should produce lower similarity."""
        steps = [
            {
                "thought": "Searching for something unrelated.",
                "action_type": "Search",
                "action_arg": "weather forecast temperature celsius",
                "observation": "Today's temperature is 22°C with partly cloudy skies.",
            }
        ]
        result = self._dispatch_retrieval_audit(self.QUESTION, steps)
        # On-topic chain (test_high_relevance_chain) should have higher sim
        on_topic_result = self._dispatch_retrieval_audit(self.QUESTION, [
            {
                "thought": "Search for battle.",
                "action_type": "Search",
                "action_arg": "Battle of Hastings 1066",
                "observation": "Fought in 1066 by William of Normandy against Harold II of England.",
            }
        ])
        # Off-topic should be lower
        assert result["mean_sim"] < on_topic_result["mean_sim"] + 0.3

    def test_multi_step_details(self):
        """Multi-step chain should produce details for each search step."""
        steps = [
            {
                "thought": "First search.",
                "action_type": "Search",
                "action_arg": "Battle of Hastings",
                "observation": "The Battle of Hastings was in 1066.",
            },
            {
                "thought": "Second search.",
                "action_type": "Search",
                "action_arg": "William the Conqueror",
                "observation": "William the Conqueror was King of England from 1066.",
            },
            {
                "thought": "Done.",
                "action_type": "Finish",
                "action_arg": "1066",
                "observation": "",
            },
        ]
        result = self._dispatch_retrieval_audit(self.QUESTION, steps)
        assert result["n_search_steps"] == 2
        assert len(result["per_step_details"]) == 2
        assert result["per_step_details"][0]["step_index"] == 0
        assert result["per_step_details"][1]["step_index"] == 1

    @pytest.mark.skipif(
        not __import__("importlib").util.find_spec("mcp"),
        reason="mcp package not installed"
    )
    def test_tool_registered_in_mcp_server(self):
        """Verify retrieval_audit appears in the MCP server's tool list."""
        from llm_guard_mcp.server import list_tools
        import asyncio
        tools = asyncio.run(list_tools())
        tool_names = [t.name for t in tools]
        assert "retrieval_audit" in tool_names

    @pytest.mark.skipif(
        not __import__("importlib").util.find_spec("mcp"),
        reason="mcp package not installed"
    )
    def test_retrieval_audit_has_required_schema_fields(self):
        """retrieval_audit tool should require question and steps."""
        from llm_guard_mcp.server import list_tools
        import asyncio
        tools = asyncio.run(list_tools())
        audit_tool = next(t for t in tools if t.name == "retrieval_audit")
        required = audit_tool.inputSchema.get("required", [])
        assert "question" in required
        assert "steps"    in required


# ── Latency benchmark helpers ─────────────────────────────────────────────────

class TestLatencyBenchmarkHelpers:
    """Unit tests for benchmark utility functions."""

    def test_percentiles_correct(self):
        """Verify percentile computation on known distribution."""
        from benchmarks.latency_benchmark import _percentiles
        # Uniform distribution [0, 100]: p50=50, p90=90, p99=99
        import numpy as np
        latencies = list(np.linspace(0, 100, 101))
        stats = _percentiles(latencies)
        assert abs(stats["p50_ms"] - 50.0) < 2.0
        assert abs(stats["p90_ms"] - 90.0) < 2.0
        assert abs(stats["p99_ms"] - 99.0) < 2.0
        assert stats["n"] == 101
        assert stats["min_ms"] <= stats["p50_ms"] <= stats["p99_ms"] <= stats["max_ms"]

    def test_single_score_returns_positive_ms(self):
        """Score should complete in positive time."""
        from llm_guard import AgentGuard
        from benchmarks.latency_benchmark import _single_score, _QUESTION, _STEPS, _ANSWER

        guard = AgentGuard()
        lat = _single_score(guard)
        assert lat > 0.0
        assert lat < 60_000.0  # sanity: under 60 seconds

    def test_benchmark_results_directory_created(self, tmp_path):
        """run_benchmark() should create output_dir if it doesn't exist."""
        from benchmarks.latency_benchmark import run_benchmark

        output_dir = tmp_path / "results"
        docs_path  = tmp_path / "latency_benchmarks.md"

        assert not output_dir.exists()

        report = run_benchmark(
            concurrency_levels=[2],
            n_warmup=1,
            n_runs=5,
            output_dir=output_dir,
            docs_path=docs_path,
        )

        assert output_dir.exists()
        json_files = list(output_dir.glob("latency_*.json"))
        assert len(json_files) == 1

    def test_benchmark_output_structure(self, tmp_path):
        """run_benchmark() should return dict with expected keys."""
        from benchmarks.latency_benchmark import run_benchmark

        report = run_benchmark(
            concurrency_levels=[2],
            n_warmup=1,
            n_runs=5,
            output_dir=tmp_path / "results",
            docs_path=tmp_path / "lb.md",
        )

        for key in ["timestamp", "version", "environment", "benchmark_params", "results"]:
            assert key in report

        assert len(report["results"]) == 1
        r = report["results"][0]
        assert r["concurrency"] == 2
        assert "p50_ms" in r
        assert "p90_ms" in r
        assert "p99_ms" in r
        assert r["n"] == 5

    def test_markdown_written(self, tmp_path):
        """run_benchmark() should write a Markdown file."""
        from benchmarks.latency_benchmark import run_benchmark

        docs_path = tmp_path / "bench.md"

        run_benchmark(
            concurrency_levels=[2],
            n_warmup=1,
            n_runs=5,
            output_dir=tmp_path / "results",
            docs_path=docs_path,
        )

        assert docs_path.exists()
        content = docs_path.read_text()
        assert "# llm-guard-kit Latency Benchmarks" in content
        assert "p50" in content
        assert "p99" in content

    def test_json_round_trip(self, tmp_path):
        """JSON output should be valid and deserializable."""
        from benchmarks.latency_benchmark import run_benchmark

        output_dir = tmp_path / "results"
        run_benchmark(
            concurrency_levels=[2],
            n_warmup=1,
            n_runs=5,
            output_dir=output_dir,
            docs_path=tmp_path / "lb.md",
        )

        json_files = list(output_dir.glob("latency_*.json"))
        assert len(json_files) == 1

        data = json.loads(json_files[0].read_text())
        assert data["results"][0]["concurrency"] == 2
        assert isinstance(data["results"][0]["p50_ms"], float)
