"""
test_cross_framework.py — Cross-agent-framework scoring validation.
=====================================================================

Validates that AgentGuard scoring produces valid, meaningful outputs for chains
generated by (or formatted as if generated by) all major agent frameworks.

Coverage matrix row: Agent framework
--------------------------------------
  ReAct native           ✓ (validated in test_claims_validation.py, exp88/89)
  OpenAI function-call   ✓ (normalised → ReAct; structural signals validated here)
  LangGraph              ✓ (same wire format as OpenAI; normalised → ReAct)
  AutoGen                ✓ (ConversableAgent format; normalised → ReAct)
  LangChain              ✓ (intermediate_steps tuples; normalised → ReAct)
  CrewAI                 ✓ (via AgentGuardCrewCallback; structural signals validated)

Important caveats — printed in every test that touches AUROC claims:
  VALIDATED: Structural output validity (risk in [0,1], tier in HIGH/MEDIUM/LOW,
             looping chains score higher than clean chains)
  NOT VALIDATED ON HELD-OUT DATA: AUROC numbers for non-ReAct frameworks.
             AUROC 0.812/0.741 was measured on Claude-generated ReAct chains
             (HotpotQA/NQ). Cross-framework AUROC may differ.
  TO VALIDATE: Run score_chain() on 200+ labeled chains per framework with
               ground-truth correctness labels and measure 5-fold CV AUROC.

Run:
    pytest tests/test_cross_framework.py -v
    pytest tests/test_cross_framework.py -v -m "not slow"
"""

import warnings
import pytest
from unittest.mock import MagicMock

warnings.filterwarnings("ignore", message=".*HF_TOKEN.*")
warnings.filterwarnings("ignore", message=".*sentence.*")

# ── Framework-specific synthetic chain fixtures ───────────────────────────────

# Clean OpenAI tool-call chain (correct answer)
OPENAI_CLEAN = [
    {
        "role": "assistant",
        "content": "I need to find who wrote Hamlet. Let me search.",
        "tool_calls": [{
            "id": "call_1",
            "type": "function",
            "function": {
                "name": "web_search",
                "arguments": '{"query": "who wrote Hamlet Shakespeare"}'
            }
        }]
    },
    {
        "role": "tool",
        "tool_call_id": "call_1",
        "content": "Hamlet is a tragedy written by William Shakespeare, believed to have been written around 1600."
    },
    {
        "role": "assistant",
        "content": "Based on the search result, Hamlet was written by William Shakespeare.",
        "tool_calls": []
    }
]

# Looping OpenAI tool-call chain (repeated searches, should score higher risk)
OPENAI_LOOP = [
    {
        "role": "assistant",
        "content": "I need to search for this.",
        "tool_calls": [{"type": "function", "function": {"name": "search", "arguments": '{"query": "x"}'}}]
    },
    {"role": "tool", "content": "No relevant results."},
    {
        "role": "assistant",
        "content": "Let me try again.",
        "tool_calls": [{"type": "function", "function": {"name": "search", "arguments": '{"query": "x"}'}}]
    },
    {"role": "tool", "content": "No relevant results."},
    {
        "role": "assistant",
        "content": "Still no results, let me try again.",
        "tool_calls": [{"type": "function", "function": {"name": "search", "arguments": '{"query": "x"}'}}]
    },
    {"role": "tool", "content": "No relevant results."},
    {
        "role": "assistant",
        "content": "I cannot find the answer.",
        "tool_calls": []
    }
]

# LangGraph format (same wire shape as OpenAI messages, with LangChain-style metadata)
LANGGRAPH_CLEAN = [
    {
        "role": "assistant",
        "content": "I'll search for information about the French capital.",
        "tool_calls": [{
            "id": "tc_1",
            "type": "function",
            "function": {"name": "tavily_search_results_json", "arguments": '{"query": "capital of France"}'}
        }]
    },
    {
        "role": "tool",
        "tool_call_id": "tc_1",
        "content": '[{"url": "...", "content": "Paris is the capital and most populous city of France."}]'
    },
    {
        "role": "assistant",
        "content": "The capital of France is Paris.",
        "tool_calls": []
    }
]

LANGGRAPH_LOOP = [
    {
        "role": "assistant",
        "content": "Searching for answer.",
        "tool_calls": [{"id": "t1", "type": "function", "function": {"name": "search", "arguments": '{"query": "unknown fact xyz"}'}}]
    },
    {"role": "tool", "tool_call_id": "t1", "content": "No results found."},
    {
        "role": "assistant",
        "content": "Trying a different search.",
        "tool_calls": [{"id": "t2", "type": "function", "function": {"name": "search", "arguments": '{"query": "unknown fact xyz"}'}}]
    },
    {"role": "tool", "tool_call_id": "t2", "content": "No results found."},
    {
        "role": "assistant",
        "content": "Unable to find information.",
        "tool_calls": []
    }
]

# AutoGen format (ConversableAgent message dicts)
AUTOGEN_CLEAN = [
    {
        "sender": "AssistantAgent",
        "role": "assistant",
        "content": (
            "Thought: I need to find the author of Hamlet.\n"
            "Action: search\n"
            "Action Input: who wrote Hamlet"
        )
    },
    {
        "sender": "UserProxyAgent",
        "role": "user",
        "content": "exitcode: 0 (execution succeeded)\nCode output: Hamlet was written by William Shakespeare."
    },
    {
        "sender": "AssistantAgent",
        "role": "assistant",
        "content": "The answer is William Shakespeare. TERMINATE"
    }
]

AUTOGEN_LOOP = [
    {
        "sender": "AssistantAgent",
        "role": "assistant",
        "content": "Thought: Searching for info.\nAction: search\nAction Input: obscure fact"
    },
    {"sender": "UserProxyAgent", "role": "user", "content": "No results."},
    {
        "sender": "AssistantAgent",
        "role": "assistant",
        "content": "Thought: Try again.\nAction: search\nAction Input: obscure fact"
    },
    {"sender": "UserProxyAgent", "role": "user", "content": "No results."},
    {
        "sender": "AssistantAgent",
        "role": "assistant",
        "content": "Thought: Still failing.\nAction: search\nAction Input: obscure fact"
    },
    {"sender": "UserProxyAgent", "role": "user", "content": "No results."},
    {"sender": "AssistantAgent", "role": "assistant", "content": "I cannot find the answer. TERMINATE"}
]

# LangChain intermediate_steps format (list of (AgentAction, observation) tuples)
LANGCHAIN_CLEAN = [
    (
        {"tool": "wikipedia", "tool_input": "William Shakespeare Hamlet", "log": "I should search Wikipedia for this."},
        "Hamlet is a tragedy by William Shakespeare, written around 1600."
    )
]

LANGCHAIN_LOOP = [
    (
        {"tool": "search", "tool_input": "obscure fact", "log": "Searching for the answer."},
        "No results found."
    ),
    (
        {"tool": "search", "tool_input": "obscure fact", "log": "Trying again."},
        "No results found."
    ),
    (
        {"tool": "search", "tool_input": "obscure fact", "log": "One more attempt."},
        "No results found."
    )
]


# ═══════════════════════════════════════════════════════════════════════════════
# Normalization correctness tests
# ═══════════════════════════════════════════════════════════════════════════════

class TestOpenAINormalization:
    """OpenAI function-call / tool-call format → canonical ReAct."""

    def test_clean_chain_normalised_keys(self):
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(OPENAI_CLEAN, agent_format="openai", warn=False)
        assert len(out) >= 1
        for s in out:
            for k in ("thought", "action_type", "action_arg", "observation"):
                assert k in s and isinstance(s[k], str)

    def test_tool_name_extracted(self):
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(OPENAI_CLEAN, agent_format="openai", warn=False)
        assert out[0]["action_type"].lower() in ("web_search", "search", "tool")

    def test_tool_result_in_observation(self):
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(OPENAI_CLEAN, agent_format="openai", warn=False)
        assert "Shakespeare" in out[0]["observation"] or "Hamlet" in out[0]["observation"]

    def test_loop_chain_normalised(self):
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(OPENAI_LOOP, agent_format="openai", warn=False)
        assert len(out) >= 2

    def test_arguments_extracted(self):
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(OPENAI_CLEAN, agent_format="openai", warn=False)
        assert "Hamlet" in out[0]["action_arg"] or "Shakespeare" in out[0]["action_arg"]


class TestLangGraphNormalization:
    """LangGraph (OpenAI message format) → canonical ReAct."""

    def test_clean_chain_normalised(self):
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(LANGGRAPH_CLEAN, agent_format="langgraph", warn=False)
        assert len(out) >= 1
        for s in out:
            for k in ("thought", "action_type", "action_arg", "observation"):
                assert k in s

    def test_observation_contains_search_result(self):
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(LANGGRAPH_CLEAN, agent_format="langgraph", warn=False)
        assert "Paris" in out[0]["observation"] or "capital" in out[0]["observation"].lower()

    def test_auto_detect_langgraph(self):
        """Auto-detect should handle LangGraph (role-based messages)."""
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(LANGGRAPH_CLEAN, agent_format="auto", warn=False)
        assert len(out) >= 1


class TestAutoGenNormalization:
    """AutoGen ConversableAgent format → canonical ReAct."""

    def test_clean_chain_normalised(self):
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(AUTOGEN_CLEAN, agent_format="autogen", warn=False)
        assert len(out) >= 1
        for s in out:
            for k in ("thought", "action_type", "action_arg", "observation"):
                assert k in s

    def test_action_extracted_from_text(self):
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(AUTOGEN_CLEAN, agent_format="autogen", warn=False)
        assert out[0]["action_type"].lower() in ("search", "action")

    def test_observation_from_user_proxy(self):
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(AUTOGEN_CLEAN, agent_format="autogen", warn=False)
        assert "Shakespeare" in out[0]["observation"] or "Hamlet" in out[0]["observation"]

    def test_auto_detect_autogen(self):
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(AUTOGEN_CLEAN, agent_format="auto", warn=False)
        assert len(out) >= 1


class TestLangChainNormalization:
    """LangChain intermediate_steps → canonical ReAct."""

    def test_clean_chain_normalised(self):
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(LANGCHAIN_CLEAN, agent_format="langchain", warn=False)
        assert len(out) == 1
        s = out[0]
        assert s["action_type"] == "wikipedia"
        assert "Shakespeare" in s["observation"] or "Hamlet" in s["observation"]

    def test_loop_chain_normalised(self):
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(LANGCHAIN_LOOP, agent_format="langchain", warn=False)
        assert len(out) == 3

    def test_log_becomes_thought(self):
        from llm_guard.features.step_normalizer import normalize_steps
        out = normalize_steps(LANGCHAIN_CLEAN, agent_format="langchain", warn=False)
        assert "Wikipedia" in out[0]["thought"] or "search" in out[0]["thought"].lower()


# ═══════════════════════════════════════════════════════════════════════════════
# Scoring validity tests (all frameworks)
# ═══════════════════════════════════════════════════════════════════════════════

class TestOpenAIScoring:
    """Scoring valid after OpenAI-format normalisation."""

    def test_risk_in_range(self):
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard(agent_format="openai")
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            r = guard.score_chain("Who wrote Hamlet?", OPENAI_CLEAN, "William Shakespeare")
        assert 0.0 <= r.risk_score <= 1.0

    def test_tier_valid(self):
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard(agent_format="openai")
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            r = guard.score_chain("Who wrote Hamlet?", OPENAI_CLEAN, "William Shakespeare")
        assert r.confidence_tier in ("HIGH", "MEDIUM", "LOW")

    def test_looping_scores_higher_than_clean(self):
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard(agent_format="openai")
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            r_clean = guard.score_chain("Who wrote Hamlet?", OPENAI_CLEAN, "William Shakespeare")
            r_loop  = guard.score_chain("Who wrote X?", OPENAI_LOOP, "unknown")
        assert r_loop.risk_score > r_clean.risk_score, (
            f"CLAIM: Looping chain ({r_loop.risk_score:.3f}) should score higher than "
            f"clean chain ({r_clean.risk_score:.3f}) for OpenAI-tools format.\n"
            f"NOTE: AUROC numbers not validated for OpenAI-tools format. "
            f"Structural signal only validated here."
        )

    def test_format_warning_emitted(self):
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard(agent_format="openai")
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            guard.score_chain("Who wrote Hamlet?", OPENAI_CLEAN, "William Shakespeare")
            format_warns = [x for x in w if "normalised" in str(x.message).lower()
                           or "openai" in str(x.message).lower()
                           or "format" in str(x.message).lower()]
            assert len(format_warns) >= 1, "Expected format warning for non-native OpenAI format"


class TestLangGraphScoring:
    """Scoring valid after LangGraph-format normalisation."""

    def test_risk_in_range(self):
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard(agent_format="langgraph")
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            r = guard.score_chain("What is the capital of France?", LANGGRAPH_CLEAN, "Paris")
        assert 0.0 <= r.risk_score <= 1.0

    def test_tier_valid(self):
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard(agent_format="langgraph")
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            r = guard.score_chain("What is the capital of France?", LANGGRAPH_CLEAN, "Paris")
        assert r.confidence_tier in ("HIGH", "MEDIUM", "LOW")

    def test_looping_scores_higher(self):
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard(agent_format="langgraph")
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            r_clean = guard.score_chain("What is the capital of France?", LANGGRAPH_CLEAN, "Paris")
            r_loop  = guard.score_chain("What is X?", LANGGRAPH_LOOP, "unknown")
        assert r_loop.risk_score > r_clean.risk_score


class TestAutoGenScoring:
    """Scoring valid after AutoGen-format normalisation."""

    def test_risk_in_range(self):
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard(agent_format="autogen")
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            r = guard.score_chain("Who wrote Hamlet?", AUTOGEN_CLEAN, "William Shakespeare")
        assert 0.0 <= r.risk_score <= 1.0

    def test_looping_scores_higher(self):
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard(agent_format="autogen")
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            r_clean = guard.score_chain("Who wrote Hamlet?", AUTOGEN_CLEAN, "William Shakespeare")
            r_loop  = guard.score_chain("Who wrote X?",  AUTOGEN_LOOP,  "unknown")
        assert r_loop.risk_score > r_clean.risk_score

    def test_needs_alert_consistent_with_risk(self):
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard(agent_format="autogen")
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            r = guard.score_chain("Q?", AUTOGEN_LOOP, "unknown")
        assert r.needs_alert == (r.risk_score >= 0.70)


class TestLangChainScoring:
    """Scoring valid after LangChain-format normalisation."""

    def test_risk_in_range(self):
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard(agent_format="langchain")
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            r = guard.score_chain("Who wrote Hamlet?", LANGCHAIN_CLEAN, "William Shakespeare")
        assert 0.0 <= r.risk_score <= 1.0

    def test_looping_scores_higher(self):
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard(agent_format="langchain")
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            r_clean = guard.score_chain("Who wrote Hamlet?", LANGCHAIN_CLEAN, "William Shakespeare")
            r_loop  = guard.score_chain("Who wrote X?", LANGCHAIN_LOOP, "unknown")
        assert r_loop.risk_score > r_clean.risk_score


class TestCrewAICallback:
    """CrewAI callback converts tool steps and scores correctly."""

    def test_crewai_on_step_and_score(self):
        from llm_guard.integrations.crewai import AgentGuardCrewCallback
        guard = MagicMock()
        guard.score_chain.return_value = MagicMock(needs_alert=False, risk_score=0.3)
        cb = AgentGuardCrewCallback(guard)
        cb.on_step({"thought": "Searching for info", "tool": "WebSearch",
                    "tool_input": "capital France", "result": "Paris is the capital"})
        task = MagicMock()
        task.raw = "Paris"
        del task.description
        cb.on_task_end(task)
        assert guard.score_chain.called

    def test_crewai_alert_on_high_risk(self):
        from llm_guard.integrations.crewai import AgentGuardCrewCallback
        guard = MagicMock()
        guard.score_chain.return_value = MagicMock(needs_alert=True, risk_score=0.85)
        alerts = []
        cb = AgentGuardCrewCallback(guard, on_alert=lambda r: alerts.append(r))
        cb.on_step({"thought": "T", "tool": "search", "tool_input": "x", "result": "nothing"})
        task = MagicMock()
        task.raw = "unknown"
        del task.description
        cb.on_task_end(task)
        assert len(alerts) == 1


# ═══════════════════════════════════════════════════════════════════════════════
# Per-call format override test
# ═══════════════════════════════════════════════════════════════════════════════

class TestPerCallFormatOverride:
    """agent_format can be overridden per score_chain() call."""

    def test_openai_override_on_react_guard(self):
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard()  # default react
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            r = guard.score_chain(
                "Who wrote Hamlet?",
                OPENAI_CLEAN,
                "William Shakespeare",
                agent_format="openai"
            )
        assert 0.0 <= r.risk_score <= 1.0

    def test_autogen_override(self):
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard()
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            r = guard.score_chain(
                "Who wrote Hamlet?",
                AUTOGEN_CLEAN,
                "William Shakespeare",
                agent_format="autogen"
            )
        assert 0.0 <= r.risk_score <= 1.0

    def test_auto_detect_works_for_all_formats(self):
        """auto format detection should work for each format."""
        from llm_guard.scoring.agent_guard import AgentGuard
        guard = AgentGuard(agent_format="auto")
        formats_and_chains = [
            (OPENAI_CLEAN,    "William Shakespeare"),
            (LANGGRAPH_CLEAN, "Paris"),
            (AUTOGEN_CLEAN,   "William Shakespeare"),
        ]
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            for steps, ans in formats_and_chains:
                r = guard.score_chain("Q?", steps, ans)
                assert 0.0 <= r.risk_score <= 1.0
                assert r.confidence_tier in ("HIGH", "MEDIUM", "LOW")


# ═══════════════════════════════════════════════════════════════════════════════
# Disclaimer test — make the limitations explicit
# ═══════════════════════════════════════════════════════════════════════════════

class TestFrameworkCoverageDisclaimer:
    """
    Documents which framework claims are validated vs. plausible-but-not-measured.

    STRUCTURAL claims (validated here):
      - Normalisation produces canonical {thought, action_type, action_arg, observation}
      - risk_score in [0, 1] for all frameworks
      - confidence_tier in HIGH/MEDIUM/LOW for all frameworks
      - Looping chains score higher risk than clean chains (directional validity)
      - UserWarning emitted for non-native formats

    AUROC claims (NOT validated for non-ReAct frameworks):
      - AUROC 0.812 within-domain: ONLY validated on Claude-generated ReAct chains, HotpotQA
      - AUROC 0.741 cross-domain: ONLY validated on Claude-generated ReAct chains, NQ
      - For OpenAI/LangGraph/AutoGen/LangChain: structural signals are model-agnostic,
        but AUROC on your specific framework+domain needs independent measurement.

    To validate for your framework:
      1. Collect 200+ chain runs with ground-truth correctness labels
      2. from llm_guard.scoring.local_verifier import LocalVerifier, extract_features
      3. Run 5-fold CV AUROC (sklearn.model_selection.cross_val_score)
      4. Compare to the 0.812 baseline
    """

    def test_coverage_matrix_is_documented(self):
        """Structural: all frameworks produce valid output. AUROC claims need per-framework validation."""
        coverage = {
            "react":     {"structural": True, "auroc_validated": True,  "auroc_value": 0.812},
            "openai":    {"structural": True, "auroc_validated": False, "auroc_value": None},
            "langgraph": {"structural": True, "auroc_validated": False, "auroc_value": None},
            "autogen":   {"structural": True, "auroc_validated": False, "auroc_value": None},
            "langchain": {"structural": True, "auroc_validated": False, "auroc_value": None},
        }
        # All frameworks validated structurally
        for fmt, info in coverage.items():
            assert info["structural"], f"{fmt} structural validation missing"

        # Only ReAct has held-out AUROC
        validated = [fmt for fmt, info in coverage.items() if info["auroc_validated"]]
        assert validated == ["react"], (
            "Only ReAct format has held-out AUROC validation. "
            "Other frameworks: structural output validated, AUROC needs per-deployment measurement."
        )
