"""Tests for pairwise evaluation (Spec 0009)."""

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 (
    EvalTask,
    EvalTrajectory,
    PairwiseJudge,
    PairwiseResult,
    PairwiseRunner,
    compare_pairwise,
)
from agentbyte.eval.base import BaseEvalTarget
from agentbyte.llm.base import BaseChatCompletionClient, BaseChatCompletionClientConfig
from agentbyte.llm.types import ChatCompletionChunk, ChatCompletionResult
from agentbyte.messages import AssistantMessage, Usage, UserMessage


# ---------------------------------------------------------------------------
# Fake client helpers
# ---------------------------------------------------------------------------


class _FakeConfig(BaseChatCompletionClientConfig):
    pass


class _FakePairwiseClient(BaseChatCompletionClient):
    """Returns pre-configured ChatCompletionResult objects in sequence."""

    def __init__(self, responses: list[ChatCompletionResult]) -> None:
        super().__init__(model="fake-pairwise", 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-pairwise", is_complete=True)

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

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


# ---------------------------------------------------------------------------
# Fake target helper
# ---------------------------------------------------------------------------


class _StaticTrajectoryTarget(BaseEvalTarget):
    """Returns a fixed trajectory regardless of the task."""

    def __init__(self, content: str, *, name: str = "static") -> None:
        super().__init__(name=name)
        self._content = content

    async def run(
        self, task: EvalTask, cancellation_token: CancellationToken | None = None
    ) -> EvalTrajectory:
        return EvalTrajectory(
            task=task,
            messages=[
                UserMessage(content=task.input, source="user"),
                AssistantMessage(content=self._content, source=self.name),
            ],
            success=True,
            usage=Usage(),
        )


# ---------------------------------------------------------------------------
# Shared trajectory fixtures
# ---------------------------------------------------------------------------


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


def _structured_response(
    winner: str, margin: float, reasoning: str
) -> ChatCompletionResult:
    from pydantic import BaseModel

    class _R(BaseModel):
        winner: str
        margin: float
        reasoning: str

    return ChatCompletionResult(
        message=AssistantMessage(content="{}", source="judge"),
        usage=Usage(tokens_input=10, tokens_output=5),
        model="fake-pairwise",
        structured_output=_R(winner=winner, margin=margin, reasoning=reasoning),
    )


# ---------------------------------------------------------------------------
# PairwiseResult model tests
# ---------------------------------------------------------------------------


def test_pairwise_result_is_tie_property() -> None:
    traj = _make_trajectory("Paris")
    result = PairwiseResult(
        task=traj.task,
        winner="tie",
        margin=0.0,
        reasoning="No difference",
        trajectory_a=traj,
        trajectory_b=traj,
    )
    assert result.is_tie is True


def test_pairwise_result_winner_a_is_not_tie() -> None:
    traj = _make_trajectory("Paris")
    result = PairwiseResult(
        task=traj.task,
        winner="a",
        margin=0.8,
        reasoning="A is better",
        trajectory_a=traj,
        trajectory_b=traj,
    )
    assert result.is_tie is False
    assert result.winner == "a"
    assert result.margin == pytest.approx(0.8)


# ---------------------------------------------------------------------------
# PairwiseJudge.compare tests
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_pairwise_judge_returns_winner_a_from_structured_response() -> None:
    client = _FakePairwiseClient([_structured_response("a", 0.9, "A is clearer")])
    judge = PairwiseJudge(client=client)

    result = await judge.compare(_make_trajectory("Paris"), _make_trajectory("The capital is Paris."))

    assert result.winner == "a"
    assert result.margin == pytest.approx(0.9)
    assert "clearer" in result.reasoning


@pytest.mark.asyncio
async def test_pairwise_judge_returns_winner_b_from_structured_response() -> None:
    client = _FakePairwiseClient([_structured_response("b", 0.6, "B is more complete")])
    judge = PairwiseJudge(client=client)

    result = await judge.compare(_make_trajectory("Paris"), _make_trajectory("The capital is Paris."))

    assert result.winner == "b"
    assert result.margin == pytest.approx(0.6)


@pytest.mark.asyncio
async def test_pairwise_judge_tie_sets_margin_to_zero() -> None:
    client = _FakePairwiseClient([_structured_response("tie", 0.3, "equivalent")])
    judge = PairwiseJudge(client=client)

    result = await judge.compare(_make_trajectory("Paris"), _make_trajectory("Paris"))

    assert result.winner == "tie"
    assert result.margin == pytest.approx(0.0)
    assert result.is_tie is True


@pytest.mark.asyncio
async def test_pairwise_judge_falls_back_to_tie_on_parse_error() -> None:
    bad_response = ChatCompletionResult(
        message=AssistantMessage(content="not-json", source="judge"),
        usage=Usage(tokens_input=5, tokens_output=2),
        model="fake-pairwise",
    )
    client = _FakePairwiseClient([bad_response])
    judge = PairwiseJudge(client=client)

    result = await judge.compare(_make_trajectory("Paris"), _make_trajectory("Paris"))

    assert result.winner == "tie"
    assert result.metadata["fallback"] is True
    assert "fallback_reason" in result.metadata


@pytest.mark.asyncio
async def test_pairwise_judge_returns_tie_on_cancellation() -> None:
    token = CancellationToken()
    token.cancel()
    client = _FakePairwiseClient([])
    judge = PairwiseJudge(client=client)

    result = await judge.compare(
        _make_trajectory("Paris"),
        _make_trajectory("Paris"),
        cancellation_token=token,
    )

    assert result.is_tie is True
    assert result.metadata["fallback"] is True


@pytest.mark.asyncio
async def test_pairwise_judge_score_raises_type_error() -> None:
    client = _FakePairwiseClient([])
    judge = PairwiseJudge(client=client)
    traj = _make_trajectory("Paris")

    with pytest.raises(TypeError, match="compare"):
        await judge.score(traj)


@pytest.mark.asyncio
async def test_pairwise_judge_clamps_margin_to_valid_range() -> None:
    """Margin values outside 0.0–1.0 from the LLM are clamped."""
    client = _FakePairwiseClient([_structured_response("a", 1.5, "very decisive")])
    judge = PairwiseJudge(client=client)

    result = await judge.compare(_make_trajectory("Paris"), _make_trajectory("Berlin"))

    assert result.margin <= 1.0


@pytest.mark.asyncio
async def test_pairwise_judge_unknown_winner_coerced_to_tie() -> None:
    """An unrecognised winner value from the LLM is coerced to tie."""
    client = _FakePairwiseClient([_structured_response("c", 0.5, "invalid winner")])
    judge = PairwiseJudge(client=client)

    result = await judge.compare(_make_trajectory("Paris"), _make_trajectory("Paris"))

    assert result.winner == "tie"


@pytest.mark.asyncio
async def test_pairwise_judge_source_filter_extracts_named_agent() -> None:
    """source_filter is applied to both trajectories before comparison."""
    task = EvalTask(name="draft", input="Write about AI", expected_output="AI draft")
    traj_a = EvalTrajectory(
        task=task,
        messages=[
            UserMessage(content="Write about AI", source="user"),
            AssistantMessage(content="AI is powerful.", source="writer"),
            AssistantMessage(content="Approved.", source="reviewer"),
        ],
        success=True,
        usage=Usage(),
    )
    traj_b = EvalTrajectory(
        task=task,
        messages=[
            UserMessage(content="Write about AI", source="user"),
            AssistantMessage(content="AI transforms industries.", source="writer"),
            AssistantMessage(content="Good work.", source="reviewer"),
        ],
        success=True,
        usage=Usage(),
    )

    client = _FakePairwiseClient([_structured_response("b", 0.7, "B writer is more informative")])
    judge = PairwiseJudge(client=client, source_filter="writer")

    result = await judge.compare(traj_a, traj_b)

    assert result.winner == "b"
    # Verify judge extracted writer messages (not reviewer)
    assert judge.extract_answer(traj_a) == "AI is powerful."
    assert judge.extract_answer(traj_b) == "AI transforms industries."


# ---------------------------------------------------------------------------
# compare_pairwise utility tests
# ---------------------------------------------------------------------------


def _make_result(winner: str) -> PairwiseResult:
    traj = _make_trajectory("Paris")
    return PairwiseResult(
        task=traj.task,
        winner=winner,  # type: ignore[arg-type]
        margin=0.5 if winner != "tie" else 0.0,
        reasoning="ok",
        trajectory_a=traj,
        trajectory_b=traj,
    )


def test_compare_pairwise_counts_wins_and_ties() -> None:
    results = [_make_result("a"), _make_result("b"), _make_result("a"), _make_result("tie")]
    summary = compare_pairwise(results, "baseline", "candidate")

    assert summary["wins_a"] == 2
    assert summary["wins_b"] == 1
    assert summary["ties"] == 1
    assert summary["total"] == 4
    assert summary["label_a"] == "baseline"
    assert summary["label_b"] == "candidate"


def test_compare_pairwise_win_rates_sum_correctly() -> None:
    results = [_make_result("a"), _make_result("b"), _make_result("a")]
    summary = compare_pairwise(results)

    assert summary["win_rate_a"] == pytest.approx(2 / 3)
    assert summary["win_rate_b"] == pytest.approx(1 / 3)


def test_compare_pairwise_empty_results_returns_zero_rates() -> None:
    summary = compare_pairwise([])
    assert summary["win_rate_a"] == 0.0
    assert summary["win_rate_b"] == 0.0
    assert summary["total"] == 0


# ---------------------------------------------------------------------------
# PairwiseRunner tests
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_pairwise_runner_returns_one_result_per_task() -> None:
    tasks = [
        EvalTask(name="t1", input="q1", expected_output="a"),
        EvalTask(name="t2", input="q2", expected_output="b"),
    ]
    client = _FakePairwiseClient(
        [_structured_response("a", 0.8, "A wins"), _structured_response("b", 0.6, "B wins")]
    )
    judge = PairwiseJudge(client=client)
    runner = PairwiseRunner(judge=judge, parallel=False)

    results = await runner.compare(
        _StaticTrajectoryTarget("answer-a", name="target-a"),
        _StaticTrajectoryTarget("answer-b", name="target-b"),
        tasks,
    )

    assert len(results) == 2
    assert results[0].winner == "a"
    assert results[1].winner == "b"


@pytest.mark.asyncio
async def test_pairwise_runner_parallel_mode_returns_all_results() -> None:
    tasks = [EvalTask(name=f"t{i}", input=f"q{i}", expected_output="x") for i in range(3)]
    client = _FakePairwiseClient(
        [_structured_response("a", 0.5, "ok")] * 3
    )
    judge = PairwiseJudge(client=client)
    runner = PairwiseRunner(judge=judge, parallel=True)

    results = await runner.compare(
        _StaticTrajectoryTarget("response-a"),
        _StaticTrajectoryTarget("response-b"),
        tasks,
    )

    assert len(results) == 3
    assert all(r.winner == "a" for r in results)


@pytest.mark.asyncio
async def test_pairwise_runner_sequential_stops_on_cancellation() -> None:
    tasks = [EvalTask(name=f"t{i}", input=f"q{i}", expected_output="x") for i in range(3)]
    token = CancellationToken()
    token.cancel()
    client = _FakePairwiseClient([])
    judge = PairwiseJudge(client=client)
    runner = PairwiseRunner(judge=judge, parallel=False)

    results = await runner.compare(
        _StaticTrajectoryTarget("a"),
        _StaticTrajectoryTarget("b"),
        tasks,
        cancellation_token=token,
    )

    assert results == []
