"""Tests for the initial AgentByte eval package (Spec 0006)."""

from __future__ import annotations

from collections.abc import AsyncGenerator
from typing import Any

import pytest

from agentbyte.cancellation_token import CancellationToken
from agentbyte.eval import (
    CompositeJudge,
    ContainsJudge,
    EvalResult,
    EvalScore,
    EvalTask,
    EvalTrajectory,
    ExactMatchJudge,
    ExpectedToolCall,
    FuzzyMatchJudge,
    LLMEvalJudge,
)
from agentbyte.llm.base import BaseChatCompletionClient, BaseChatCompletionClientConfig
from agentbyte.llm.types import ChatCompletionChunk, ChatCompletionResult
from agentbyte.messages import AssistantMessage, SystemMessage, ToolMessage, Usage, UserMessage


class FakeClientConfig(BaseChatCompletionClientConfig):
    """Config stub for fake eval clients."""


class FakeJudgeClient(BaseChatCompletionClient):
    """Simple fake chat client for LLM-judge tests."""

    def __init__(self, responses: list[ChatCompletionResult]) -> None:
        super().__init__(model="fake-judge", client=object())
        self._responses = responses
        self._index = 0

    async def create(
        self,
        messages: list,
        tools: list[dict[str, Any]] | None = None,
        output_format: type | None = None,
        **kwargs: Any,
    ) -> ChatCompletionResult:
        response = self._responses[self._index]
        self._index += 1
        return response

    async def create_stream(
        self,
        messages: list,
        tools: list[dict[str, Any]] | None = None,
        output_format: type | None = None,
        **kwargs: Any,
    ) -> AsyncGenerator[ChatCompletionChunk, None]:
        yield ChatCompletionChunk(content="", model="fake-judge", is_complete=True)

    def _to_config(self) -> BaseChatCompletionClientConfig:
        return FakeClientConfig(model=self.model, config=self.config)

    @classmethod
    def _from_config(cls, config: BaseChatCompletionClientConfig) -> "FakeJudgeClient":
        return cls(responses=[])


class StaticScoreJudge(ExactMatchJudge):
    """Test helper judge that returns a fixed score payload."""

    def __init__(
        self,
        *,
        name: str,
        overall: float,
        dimensions: dict[str, float],
        reasoning: dict[str, str] | None = None,
    ) -> None:
        super().__init__(name=name)
        self._overall = overall
        self._dimensions = dimensions
        self._reasoning = reasoning or {
            dimension: f"{name} scored {score}"
            for dimension, score in dimensions.items()
        }

    async def score(
        self,
        trajectory: EvalTrajectory,
        criteria: list[str] | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> EvalScore:
        return EvalScore(
            overall=self._overall,
            dimensions=self._dimensions,
            reasoning=self._reasoning,
            trajectory=trajectory,
            metadata={"judge": self.name},
        )


def _build_success_trajectory(
    *,
    answer: str = "Paris",
    expected_output: str | None = "Paris",
) -> EvalTrajectory:
    task = EvalTask(
        name="capital",
        input="What is the capital of France?",
        expected_output=expected_output,
    )
    return EvalTrajectory(
        task=task,
        messages=[
            UserMessage(content=task.input, source="user"),
            AssistantMessage(content=answer, source="assistant"),
        ],
        success=True,
        usage=Usage(tokens_input=5, tokens_output=2),
    )


def test_eval_task_supports_expected_tool_calls() -> None:
    task = EvalTask(
        name="weather",
        input="Weather in Paris",
        expected_tool_calls=[ExpectedToolCall(name="get_weather", arguments={"location": "Paris"})],
    )

    assert task.expected_tool_calls is not None
    assert task.expected_tool_calls[0].name == "get_weather"


def test_eval_score_helpers_return_conversation_views() -> None:
    trajectory = _build_success_trajectory(answer="Paris")
    score = EvalScore(
        overall=10.0,
        dimensions={"accuracy": 10.0},
        reasoning={"accuracy": "Correct"},
        trajectory=trajectory,
    )

    assert score.get_final_response() == "Paris"
    assert "What is the capital of France?" in score.get_full_conversation()
    assert "Paris" in score.get_full_conversation()


def test_eval_result_bundles_task_trajectory_and_score() -> None:
    trajectory = _build_success_trajectory()
    score = EvalScore(
        overall=10.0,
        dimensions={"accuracy": 10.0},
        reasoning={"accuracy": "Correct"},
        trajectory=trajectory,
    )
    result = EvalResult(
        task=trajectory.task,
        trajectory=trajectory,
        score=score,
        judge_name="ExactMatch",
    )

    assert result.judge_name == "ExactMatch"
    assert result.score.overall == 10.0


@pytest.mark.asyncio
async def test_exact_match_judge_scores_exact_match() -> None:
    judge = ExactMatchJudge()
    score = await judge.score(_build_success_trajectory())

    assert score.overall == 10.0
    assert score.dimensions["accuracy"] == 10.0


@pytest.mark.asyncio
async def test_contains_judge_scores_substring_match() -> None:
    judge = ContainsJudge()
    score = await judge.score(
        _build_success_trajectory(answer="The capital is Paris.")
    )

    assert score.overall == 10.0


@pytest.mark.asyncio
async def test_fuzzy_match_judge_scores_partial_similarity() -> None:
    judge = FuzzyMatchJudge(threshold=0.8)
    score = await judge.score(
        _build_success_trajectory(answer="Pariss")
    )

    assert 0.0 < score.overall <= 10.0
    assert "ratio" in score.reasoning["accuracy"]


@pytest.mark.asyncio
async def test_reference_judges_return_zero_for_failed_trajectory() -> None:
    trajectory = EvalTrajectory(
        task=EvalTask(name="failed", input="Broken", expected_output="ok"),
        messages=[],
        success=False,
        error="boom",
        usage=Usage(),
    )

    score = await ExactMatchJudge().score(trajectory)

    assert score.overall == 0.0
    assert score.reasoning["accuracy"] == "boom"


@pytest.mark.asyncio
async def test_base_answer_extraction_last_assistant_skips_tool_messages() -> None:
    judge = ExactMatchJudge(answer_strategy="last_assistant")
    trajectory = EvalTrajectory(
        task=EvalTask(name="math", input="2+2", expected_output="4"),
        messages=[
            UserMessage(content="2+2", source="user"),
            AssistantMessage(content="Let me calculate", source="assistant"),
            ToolMessage(
                content="4",
                source="tool",
                tool_call_id="call-1",
                tool_name="calculator",
                success=True,
            ),
            AssistantMessage(content="4", source="assistant"),
        ],
        success=True,
        usage=Usage(),
    )

    assert judge.extract_answer(trajectory) == "4"


def _multi_agent_trajectory() -> EvalTrajectory:
    """Flat trajectory from a Researcher→Writer→Reviewer orchestrator run."""
    task = EvalTask(name="draft", input="Write about AI", expected_output="AI draft")
    return EvalTrajectory(
        task=task,
        messages=[
            UserMessage(content="Write about AI", source="user"),
            AssistantMessage(content="First draft: AI is powerful.", source="writer"),
            AssistantMessage(content="Second draft: AI transforms industries.", source="writer"),
            AssistantMessage(content="Approved. Looks good.", source="reviewer"),
        ],
        success=True,
        usage=Usage(),
    )


def test_source_filter_last_non_empty_returns_last_matching_agent_message() -> None:
    judge = ExactMatchJudge(answer_strategy="last_non_empty", source_filter="writer")
    answer = judge.extract_answer(_multi_agent_trajectory())

    assert answer == "Second draft: AI transforms industries."


def test_source_filter_all_assistant_returns_all_matching_agent_messages() -> None:
    judge = ExactMatchJudge(answer_strategy="all_assistant", source_filter="writer")
    answer = judge.extract_answer(_multi_agent_trajectory())

    assert "First draft" in answer
    assert "Second draft" in answer
    assert "Approved" not in answer


def test_source_filter_last_assistant_ignores_other_agent_messages() -> None:
    judge = ExactMatchJudge(answer_strategy="last_assistant", source_filter="reviewer")
    answer = judge.extract_answer(_multi_agent_trajectory())

    assert answer == "Approved. Looks good."
    assert "draft" not in answer.lower()


def test_source_filter_unknown_source_returns_empty_string() -> None:
    judge = ExactMatchJudge(source_filter="nonexistent")
    answer = judge.extract_answer(_multi_agent_trajectory())

    assert answer == ""


@pytest.mark.asyncio
async def test_contains_judge_with_source_filter_scores_target_agent() -> None:
    """ContainsJudge(source_filter="writer") matches against the writer's output,
    not the reviewer's — even though the reviewer's message is last overall."""
    task = EvalTask(name="draft", input="Write about AI", expected_output="transforms industries")
    trajectory = EvalTrajectory(
        task=task,
        messages=_multi_agent_trajectory().messages,
        success=True,
        usage=Usage(),
    )
    judge = ContainsJudge(source_filter="writer")

    score = await judge.score(trajectory)

    assert score.overall == 10.0


@pytest.mark.asyncio
async def test_llm_eval_judge_uses_structured_response_when_available() -> None:
    client = FakeJudgeClient(
        responses=[
            ChatCompletionResult(
                message=AssistantMessage(content='{"overall": 9.0}', source="judge"),
                usage=Usage(tokens_input=10, tokens_output=5),
                model="fake-judge",
                structured_output={
                    "overall": 9.0,
                    "dimensions": {"accuracy": 8.0, "helpfulness": 10.0},
                    "reasoning": {
                        "accuracy": "Accurate",
                        "helpfulness": "Helpful",
                    },
                },
            )
        ]
    )
    judge = LLMEvalJudge(client=client, default_criteria=["accuracy", "helpfulness"])

    score = await judge.score(_build_success_trajectory(), criteria=None)

    assert score.overall == 9.0
    assert score.dimensions == {"accuracy": 8.0, "helpfulness": 10.0}


@pytest.mark.asyncio
async def test_llm_eval_judge_returns_neutral_score_on_parse_error() -> None:
    client = FakeJudgeClient(
        responses=[
            ChatCompletionResult(
                message=AssistantMessage(content="not-json", source="judge"),
                usage=Usage(tokens_input=10, tokens_output=5),
                model="fake-judge",
            )
        ]
    )
    judge = LLMEvalJudge(client=client)

    score = await judge.score(_build_success_trajectory())

    assert score.overall == 5.0
    assert all(value == 5.0 for value in score.dimensions.values())
    assert score.metadata["fallback"] is True


@pytest.mark.asyncio
async def test_llm_eval_judge_returns_neutral_score_when_cancelled() -> None:
    token = CancellationToken()
    token.cancel()
    client = FakeJudgeClient(responses=[])
    judge = LLMEvalJudge(client=client)

    score = await judge.score(_build_success_trajectory(), cancellation_token=token)

    assert score.overall == 5.0
    assert score.metadata["fallback"] is True


@pytest.mark.asyncio
async def test_composite_judge_combines_sub_judge_scores() -> None:
    trajectory = _build_success_trajectory(answer="The answer is Paris.")
    judge = CompositeJudge(
        judges=[
            (ContainsJudge(name="contains"), 1.0),
            (FuzzyMatchJudge(name="fuzzy"), 3.0),
        ]
    )

    score = await judge.score(trajectory)

    assert score.overall > 0.0
    assert "accuracy (contains)" in score.reasoning
    assert "sub_judges" in score.metadata


@pytest.mark.asyncio
async def test_composite_judge_overall_uses_judge_weights_across_distinct_dimensions() -> None:
    trajectory = _build_success_trajectory()
    judge = CompositeJudge(
        judges=[
            (
                StaticScoreJudge(
                    name="accuracy_judge",
                    overall=10.0,
                    dimensions={"accuracy": 10.0},
                ),
                0.2,
            ),
            (
                StaticScoreJudge(
                    name="helpfulness_judge",
                    overall=4.0,
                    dimensions={"helpfulness": 4.0},
                ),
                0.8,
            ),
        ],
        normalize_weights=False,
    )

    score = await judge.score(trajectory)

    assert score.overall == pytest.approx(5.2)
    assert score.dimensions == {"accuracy": 10.0, "helpfulness": 4.0}
    assert "accuracy (accuracy_judge)" in score.reasoning
    assert "helpfulness (helpfulness_judge)" in score.reasoning


def test_composite_judge_rejects_invalid_non_normalized_weights() -> None:
    with pytest.raises(ValueError):
        CompositeJudge(
            judges=[(ExactMatchJudge(), 1.0), (ContainsJudge(), 1.0)],
            normalize_weights=False,
        )


def test_answer_strategy_validation_rejects_unknown_value() -> None:
    with pytest.raises(ValueError):
        ExactMatchJudge(answer_strategy="unknown")  # type: ignore[arg-type]


def test_source_filter_works_on_single_agent_trajectory() -> None:
    """source_filter on a single-agent trajectory (agent name as source)."""
    task = EvalTask(name="q", input="hello", expected_output="world")
    trajectory = EvalTrajectory(
        task=task,
        messages=[
            UserMessage(content="hello", source="user"),
            AssistantMessage(content="world", source="my-agent"),
        ],
        success=True,
        usage=Usage(),
    )
    judge = ExactMatchJudge(source_filter="my-agent")
    assert judge.extract_answer(trajectory) == "world"


def test_source_filter_works_on_raw_model_trajectory_and_is_case_sensitive() -> None:
    """source_filter on a ModelEvalTarget-style trajectory (model name as source).

    Also verifies that matching is exact and case-sensitive: 'GPT-4' does not
    match 'gpt-4', so a mismatched filter returns an empty string.
    """
    task = EvalTask(name="q", input="hello", expected_output="world")
    trajectory = EvalTrajectory(
        task=task,
        messages=[
            SystemMessage(content="Be helpful", source="system"),
            UserMessage(content="hello", source="user"),
            AssistantMessage(content="world", source="gpt-4"),
        ],
        success=True,
        usage=Usage(),
    )
    judge = ExactMatchJudge(source_filter="gpt-4")
    assert judge.extract_answer(trajectory) == "world"

    judge_wrong_case = ExactMatchJudge(source_filter="GPT-4")
    assert judge_wrong_case.extract_answer(trajectory) == ""


@pytest.mark.asyncio
async def test_composite_judge_weighted_average_math() -> None:
    """CompositeJudge overall is the weighted average of sub-judge scores.

    answer = "The answer is Paris.", expected_output = "Paris"
    ContainsJudge  (weight 3, normalised 0.75): "Paris" in answer → 10.0
    ExactMatchJudge (weight 1, normalised 0.25): "Paris" != answer →  0.0
    expected overall: 10.0 * 0.75 + 0.0 * 0.25 = 7.5
    """
    trajectory = _build_success_trajectory(
        answer="The answer is Paris.",
        expected_output="Paris",
    )
    judge = CompositeJudge(
        judges=[
            (ContainsJudge(name="contains"), 3.0),
            (ExactMatchJudge(name="exact"), 1.0),
        ]
    )
    score = await judge.score(trajectory)

    assert pytest.approx(score.overall, abs=0.01) == 7.5
    assert score.dimensions["accuracy"] == pytest.approx(7.5, abs=0.01)
