"""Tests for Spec 0008 Phase 1 runner and targets."""

from __future__ import annotations

from collections.abc import AsyncGenerator
from typing import Any

import pytest

from agentbyte.agents import BaseAgent
from agentbyte.agents.types import AgentEvent, AgentResponse
from agentbyte.cancellation_token import CancellationToken
from agentbyte.context import AgentContext
from agentbyte.eval import (
    AgentEvalTarget,
    BaseEvalJudge,
    EvalRunner,
    EvalScore,
    EvalTask,
    EvalTrajectory,
    ExactMatchJudge,
    ModelEvalTarget,
    OrchestratorEvalTarget,
    compare_configurations,
)
from agentbyte.llm.base import BaseChatCompletionClient, BaseChatCompletionClientConfig
from agentbyte.llm.types import ChatCompletionChunk, ChatCompletionResult
from agentbyte.messages import AssistantMessage, Message, Usage
from agentbyte.orchestration import BaseOrchestrator
from agentbyte.termination import MaxMessageTermination
from agentbyte.types import OrchestrationEvent, OrchestrationResponse, StopMessage


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


class FakeModelClient(BaseChatCompletionClient):
    """Simple fake chat client for target tests."""

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

    async def create(
        self,
        messages: list[Message],
        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[Message],
        tools: list[dict[str, Any]] | None = None,
        output_format: type | None = None,
        **kwargs: Any,
    ) -> AsyncGenerator[ChatCompletionChunk, None]:
        yield ChatCompletionChunk(content="", model="fake-model", is_complete=True)

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

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


class FakeAgent(BaseAgent):
    """Minimal fake agent for eval target tests."""

    def __init__(self, *, response_text: str = "done", fail: bool = False) -> None:
        super().__init__(
            name="fake-agent",
            description="fake",
            instructions="help",
            model_client=FakeModelClient(responses=[]),
        )
        self.response_text = response_text
        self.fail = fail

    async def run(
        self,
        task: str | Message | list[Message] | None = None,
        context: AgentContext | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> AgentResponse:
        if self.fail:
            raise RuntimeError("agent failed")
        ctx = context or AgentContext()
        for message in self._convert_task_to_messages(task or ""):
            ctx.messages.append(message)
        ctx.messages.append(
            AssistantMessage(content=self.response_text, source=self.name)
        )
        return AgentResponse(
            context=ctx,
            source=self.name,
            finish_reason="stop",
            usage=Usage(duration_ms=7, llm_calls=1),
        )

    async def run_stream(
        self,
        task: str | Message | list[Message] | None = None,
        context: AgentContext | None = None,
        cancellation_token: CancellationToken | None = None,
        verbose: bool = False,
        stream_tokens: bool = False,
    ) -> AsyncGenerator[Message | AgentEvent | AgentResponse | ChatCompletionChunk, None]:
        yield await self.run(task, context=context, cancellation_token=cancellation_token)


class FakeOrchestrator(BaseOrchestrator):
    """Minimal orchestrator used to test the eval target wrapper."""

    def __init__(self, *, fail: bool = False) -> None:
        super().__init__(
            agents=[FakeAgent()],
            termination=MaxMessageTermination(max_messages=5),
        )
        self.fail = fail

    async def select_next_agent(self) -> BaseAgent:
        return self.agents[0]

    async def prepare_context_for_agent(self, agent: BaseAgent) -> Message | list[Message]:
        raise NotImplementedError

    async def update_shared_state(self, result: AgentResponse) -> None:
        raise NotImplementedError

    async def run(
        self,
        task: str | Message | list[Message],
        cancellation_token: CancellationToken | None = None,
        context: AgentContext | None = None,
    ) -> OrchestrationResponse:
        if self.fail:
            raise RuntimeError("orchestrator failed")
        return OrchestrationResponse(
            messages=[AssistantMessage(content="team result", source="worker")],
            final_result="team result",
            usage=Usage(duration_ms=11, llm_calls=2),
            stop_message=StopMessage(content="Done", source="TestStop"),
            pattern_metadata={"pattern": "fake", "iterations_completed": 3},
        )

    async def run_stream(
        self,
        task: str | Message | list[Message],
        cancellation_token: CancellationToken | None = None,
        verbose: bool = False,
        context: AgentContext | None = None,
    ) -> AsyncGenerator[Message | OrchestrationEvent | OrchestrationResponse, None]:
        yield await self.run(task, cancellation_token=cancellation_token, context=context)


class FailingJudge(BaseEvalJudge):
    """Judge that raises to test runner isolation."""

    def __init__(self) -> None:
        super().__init__(name="FailingJudge")

    async def score(
        self,
        trajectory: EvalTrajectory,
        criteria: list[str] | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> EvalScore:
        raise RuntimeError("judge exploded")


@pytest.mark.asyncio
async def test_model_eval_target_wraps_chat_client_result() -> None:
    target = ModelEvalTarget(
        client=FakeModelClient(
            responses=[
                ChatCompletionResult(
                    message=AssistantMessage(content="Paris", source="model"),
                    usage=Usage(duration_ms=3, llm_calls=1),
                    model="fake-model",
                    finish_reason="stop",
                )
            ]
        ),
        system_message="Be concise",
    )

    trajectory = await target.run(
        EvalTask(name="capital", input="Capital of France?", expected_output="Paris")
    )

    assert trajectory.success is True
    assert trajectory.messages[-1].content == "Paris"
    assert trajectory.metadata["target_type"] == "model"
    assert trajectory.metadata["finish_reason"] == "stop"


@pytest.mark.asyncio
async def test_agent_eval_target_wraps_agent_response() -> None:
    target = AgentEvalTarget(FakeAgent(response_text="done"))

    trajectory = await target.run(EvalTask(name="x", input="hello"))

    assert trajectory.success is True
    assert trajectory.messages[-1].content == "done"
    assert trajectory.metadata["target_type"] == "agent"


@pytest.mark.asyncio
async def test_orchestrator_eval_target_wraps_orchestration_response() -> None:
    target = OrchestratorEvalTarget(FakeOrchestrator())

    trajectory = await target.run(EvalTask(name="x", input="hello"))

    assert trajectory.success is True
    assert trajectory.metadata["target_type"] == "orchestrator"
    assert trajectory.metadata["pattern"] == "fake"
    assert trajectory.metadata["iterations"] == 3
    assert trajectory.metadata["stop_reason"] == "Done"


@pytest.mark.asyncio
async def test_targets_return_failed_trajectory_instead_of_raising() -> None:
    agent_trajectory = await AgentEvalTarget(FakeAgent(fail=True)).run(
        EvalTask(name="x", input="hello")
    )
    orchestrator_trajectory = await OrchestratorEvalTarget(
        FakeOrchestrator(fail=True)
    ).run(EvalTask(name="x", input="hello"))

    assert agent_trajectory.success is False
    assert "agent failed" in (agent_trajectory.error or "")
    assert orchestrator_trajectory.success is False
    assert "orchestrator failed" in (orchestrator_trajectory.error or "")


@pytest.mark.asyncio
async def test_eval_runner_parallel_returns_scores_for_all_tasks() -> None:
    runner = EvalRunner(judge=ExactMatchJudge(), parallel=True)
    scores = await runner.evaluate(
        AgentEvalTarget(FakeAgent(response_text="Paris")),
        [
            EvalTask(name="capital-1", input="A", expected_output="Paris"),
            EvalTask(name="capital-2", input="B", expected_output="Paris"),
        ],
    )

    assert len(scores) == 2
    assert all(score.overall == 10.0 for score in scores)


@pytest.mark.asyncio
async def test_eval_runner_sequential_stops_early_on_cancellation() -> None:
    token = CancellationToken()
    token.cancel()
    runner = EvalRunner(judge=ExactMatchJudge(), parallel=False)

    scores = await runner.evaluate(
        AgentEvalTarget(FakeAgent(response_text="Paris")),
        [
            EvalTask(name="capital-1", input="A", expected_output="Paris"),
            EvalTask(name="capital-2", input="B", expected_output="Paris"),
        ],
        cancellation_token=token,
    )

    assert scores == []


@pytest.mark.asyncio
async def test_eval_runner_returns_zero_score_when_judge_raises() -> None:
    runner = EvalRunner(judge=FailingJudge(), parallel=True)

    scores = await runner.evaluate(
        AgentEvalTarget(FakeAgent(response_text="Paris")),
        [EvalTask(name="capital", input="A", expected_output="Paris")],
    )

    assert len(scores) == 1
    assert scores[0].overall == 0.0
    # When the judge fails the actual run trajectory is preserved (agent succeeded)
    assert scores[0].trajectory is not None
    assert scores[0].trajectory.success is True
    assert "judge exploded" in (scores[0].metadata.get("error") or "")
    assert scores[0].metadata.get("judge_failed") is True


def test_compare_configurations_computes_averages_and_improvement() -> None:
    baseline = [
        EvalScore(overall=8.0, trajectory=None),
        EvalScore(overall=10.0, trajectory=None),
    ]
    candidate = [
        EvalScore(overall=9.0, trajectory=None),
        EvalScore(overall=10.0, trajectory=None),
    ]

    result = compare_configurations(
        baseline,
        candidate,
        label_a="baseline",
        label_b="candidate",
    )

    assert result["label_a"] == "baseline"
    assert result["label_b"] == "candidate"
    assert result["avg_a"] == 9.0
    assert result["avg_b"] == 9.5
    assert pytest.approx(result["improvement_pct"], rel=1e-6) == (
        (9.5 - 9.0) / 9.0
    ) * 100.0
