from __future__ import annotations

from collections.abc import AsyncGenerator
from typing import Any, Optional, Union

import pytest

from agentbyte.agents.base import BaseAgent
from agentbyte.agents.types import AgentResponse, Usage
from agentbyte.context import AgentContext
from agentbyte.llm.base import BaseChatCompletionClient, BaseChatCompletionClientConfig
from agentbyte.llm.types import ChatCompletionChunk, ChatCompletionResult
from agentbyte.messages import AssistantMessage, Message, UserMessage
from agentbyte.orchestration import (
    AIOrchestrator,
    AgentHistoryPolicy,
    AgentSelection,
    HistoryContextPolicy,
    SelectorContextPolicy,
)
from agentbyte.termination.max_message import MaxMessageTermination
from agentbyte.types import AgentSelectionEvent, OrchestrationResponse


class FakeSelectorConfig(BaseChatCompletionClientConfig):
    pass


class FakeSelectorClient(BaseChatCompletionClient):
    def __init__(
        self,
        responses: list[ChatCompletionResult],
        fail_on_calls: Optional[set[int]] = None,
    ) -> None:
        super().__init__(model="selector-fake", client=object())
        self._responses = responses
        self._index = 0
        self._fail_on_calls = fail_on_calls or set()
        self.prompts: list[str] = []

    async def create(
        self,
        messages: list[Message],
        tools: Optional[list[dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> ChatCompletionResult:
        del tools, output_format, kwargs
        self._index += 1
        self.prompts.append(messages[0].content if messages else "")
        if self._index in self._fail_on_calls:
            raise RuntimeError("selector failed")
        return self._responses[self._index - 1]

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

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

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


class MockAgent(BaseAgent):
    def __init__(self, name: str, response_text: str):
        self.name = name
        self.description = f"Description of {name}"
        self.response_text = response_text
        self.tools = []
        self.received_tasks: list[Union[UserMessage, list[Message]]] = []

    async def run(
        self,
        task: Optional[Union[UserMessage, list[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[Any] = None,
    ) -> AgentResponse:
        del task, context, cancellation_token
        raise NotImplementedError("Use run_stream")

    async def run_stream(
        self,
        task: Optional[Union[UserMessage, list[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[Any] = None,
        verbose: bool = False,
        stream_tokens: bool = False,
    ) -> AsyncGenerator[Union[Message, Any, AgentResponse], None]:
        del cancellation_token, verbose, stream_tokens
        ctx = context or AgentContext()
        if task is not None:
            self.received_tasks.append(task)

        if isinstance(task, UserMessage):
            yield task
            if not ctx.messages:
                ctx.messages.append(task)
        elif isinstance(task, list):
            for message in task:
                yield message
                if not ctx.messages:
                    ctx.messages.append(message)

        message = AssistantMessage(content=self.response_text, source=self.name)
        yield message

        ctx.messages.append(message)
        yield AgentResponse(
            source=self.name,
            context=ctx,
            usage=Usage(tokens_input=5, tokens_output=5, llm_calls=1),
            finish_reason="stop",
        )


def make_selection_result(
    agent_name: str,
    reasoning: str,
    tokens_input: int = 3,
    tokens_output: int = 1,
) -> ChatCompletionResult:
    return ChatCompletionResult(
        message=AssistantMessage(content="selection", source="selector"),
        usage=Usage(
            tokens_input=tokens_input,
            tokens_output=tokens_output,
            llm_calls=1,
        ),
        model="selector-fake",
        structured_output=AgentSelection(selected_agent=agent_name, reasoning=reasoning),
    )


@pytest.mark.asyncio
class TestAIOrchestrator:
    async def test_ai_selection_chooses_requested_agent(self):
        a1 = MockAgent("researcher", "research output")
        a2 = MockAgent("writer", "writer output")
        selector = FakeSelectorClient(
            responses=[
                make_selection_result("writer", "Need synthesis first"),
                make_selection_result("researcher", "Need additional facts"),
            ]
        )
        orchestrator = AIOrchestrator(
            agents=[a1, a2],
            termination=MaxMessageTermination(3),
            model_client=selector,
            max_iterations=3,
        )

        response = await orchestrator.run(UserMessage(content="start", source="user"))

        assert isinstance(response, OrchestrationResponse)
        assert response.messages[1].source == "writer"
        assert response.messages[2].source == "researcher"
        assert response.pattern_metadata["selector_calls"] == 2

    async def test_unknown_agent_falls_back_to_round_robin(self):
        a1 = MockAgent("researcher", "research output")
        a2 = MockAgent("writer", "writer output")
        selector = FakeSelectorClient(
            responses=[make_selection_result("ghost", "Wrong agent name")]
        )
        orchestrator = AIOrchestrator(
            agents=[a1, a2],
            termination=MaxMessageTermination(2),
            model_client=selector,
        )

        selection_events: list[AgentSelectionEvent] = []
        async for item in orchestrator.run_stream(
            UserMessage(content="start", source="user"),
            verbose=True,
        ):
            if isinstance(item, AgentSelectionEvent):
                selection_events.append(item)

        assert selection_events
        assert "Fallback round-robin selected 'researcher'" in selection_events[0].selection_reason

    async def test_selector_error_falls_back_to_round_robin(self):
        a1 = MockAgent("researcher", "research output")
        a2 = MockAgent("writer", "writer output")
        selector = FakeSelectorClient(
            responses=[make_selection_result("writer", "unused")],
            fail_on_calls={1},
        )
        orchestrator = AIOrchestrator(
            agents=[a1, a2],
            termination=MaxMessageTermination(2),
            model_client=selector,
        )

        response = await orchestrator.run(UserMessage(content="start", source="user"))

        assert isinstance(response, OrchestrationResponse)
        assert response.messages[1].source == "researcher"

    async def test_usage_includes_selector_calls(self):
        a1 = MockAgent("researcher", "research output")
        selector = FakeSelectorClient(
            responses=[make_selection_result("researcher", "Only agent", 7, 2)]
        )
        orchestrator = AIOrchestrator(
            agents=[a1],
            termination=MaxMessageTermination(2),
            model_client=selector,
        )

        response = await orchestrator.run(UserMessage(content="start", source="user"))

        assert response.usage.tokens_input == 12
        assert response.usage.tokens_output == 7
        assert response.usage.llm_calls == 2

    async def test_selection_supports_case_insensitive_agent_match(self):
        a1 = MockAgent("researcher", "research output")
        a2 = MockAgent("writer", "writer output")
        selector = FakeSelectorClient(
            responses=[make_selection_result("WRITER", "Need writing")]
        )
        orchestrator = AIOrchestrator(
            agents=[a1, a2],
            termination=MaxMessageTermination(2),
            model_client=selector,
        )

        response = await orchestrator.run(UserMessage(content="start", source="user"))

        assert response.messages[1].source == "writer"

    async def test_selector_policy_is_used_for_prompt_rendering(self):
        researcher = MockAgent("researcher", "research output")
        writer = MockAgent("writer", "writer output")
        selector = FakeSelectorClient(
            responses=[make_selection_result("writer", "Need writing")]
        )
        orchestrator = AIOrchestrator(
            agents=[researcher, writer],
            termination=MaxMessageTermination(2),
            model_client=selector,
            selector_policy=SelectorContextPolicy(
                window="recent",
                message_limit=1,
                format="json",
            ),
        )
        orchestrator.shared_messages = [
            UserMessage(content="first", source="user"),
            AssistantMessage(content="second", source="researcher"),
        ]

        prompt = orchestrator._build_selection_prompt()

        assert '"conversation"' in prompt
        assert '"content": "second"' in prompt
        assert '"content": "first"' not in prompt

    async def test_worker_history_policy_is_independent_from_selector_policy(self):
        researcher = MockAgent("researcher", "research output")
        writer = MockAgent("writer", "writer output")
        selector = FakeSelectorClient(
            responses=[make_selection_result("writer", "Need writing")]
        )
        orchestrator = AIOrchestrator(
            agents=[researcher, writer],
            termination=MaxMessageTermination(2),
            model_client=selector,
            history_policy=HistoryContextPolicy(
                default=AgentHistoryPolicy(
                    window="recent",
                    message_limit=1,
                    format="messages",
                )
            ),
            selector_policy=SelectorContextPolicy(
                window="recent",
                message_limit=2,
                format="summary",
            ),
        )
        orchestrator.shared_messages = [
            UserMessage(content="first", source="user"),
            AssistantMessage(content="second", source="researcher"),
            AssistantMessage(content="third", source="writer"),
        ]

        context = await orchestrator.prepare_context_for_agent(writer)

        assert isinstance(context, list)
        assert [message.content for message in context] == ["third"]
        assert "second" in orchestrator._build_selection_prompt()
        assert "first" not in orchestrator._build_selection_prompt()

    async def test_selector_markdown_format_is_rendered(self):
        researcher = MockAgent("researcher", "research output")
        selector = FakeSelectorClient(
            responses=[make_selection_result("researcher", "Only agent")]
        )
        orchestrator = AIOrchestrator(
            agents=[researcher],
            termination=MaxMessageTermination(2),
            model_client=selector,
            selector_policy=SelectorContextPolicy(format="markdown"),
        )
        orchestrator.shared_messages = [
            UserMessage(content="start", source="user"),
            AssistantMessage(content="work", source="researcher"),
        ]

        context = orchestrator._format_conversation_for_selection()

        assert "# Recent Conversation" in context
        assert "## researcher (assistant)" in context
