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.cancellation_token import CancellationToken
from agentbyte.context import AgentContext
from agentbyte.messages import AssistantMessage, Message, UserMessage
from agentbyte.orchestration import AgentHistoryPolicy, HistoryContextPolicy
from agentbyte.orchestration.round_robin import RoundRobinOrchestrator
from agentbyte.termination.max_message import MaxMessageTermination


class MockAgent(BaseAgent):
    def __init__(self, name: str, responses: list[str]):
        self.name = name
        self.description = f"Description of {name}"
        self.responses = responses
        self._resp_idx = 0
        self.tools = []

    async def run(
        self,
        task: Optional[Union[UserMessage, list[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[CancellationToken] = None,
    ) -> AgentResponse:
        del cancellation_token
        resp_text = self.responses[self._resp_idx % len(self.responses)]
        self._resp_idx += 1
        ctx = context or AgentContext()
        if task and not ctx.messages:
            if isinstance(task, UserMessage):
                ctx.messages.append(task)
            else:
                ctx.messages.extend(task)
        msg = AssistantMessage(content=resp_text, source=self.name)
        ctx.messages.append(msg)
        return AgentResponse(source=self.name, context=ctx, finish_reason="stop")

    async def run_stream(
        self,
        task: Optional[Union[UserMessage, list[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[CancellationToken] = None,
        verbose: bool = False,
        stream_tokens: bool = False,
    ) -> AsyncGenerator[Union[Message, Any, AgentResponse], None]:
        del cancellation_token, verbose, stream_tokens
        resp_text = self.responses[self._resp_idx % len(self.responses)]
        self._resp_idx += 1

        ctx = context or AgentContext()
        if task:
            if isinstance(task, UserMessage):
                yield task
                if not ctx.messages:
                    ctx.messages.append(task)
            else:
                for message in task:
                    yield message
                    if not ctx.messages:
                        ctx.messages.append(message)

        msg = AssistantMessage(content=resp_text, source=self.name)
        yield msg
        ctx.messages.append(msg)
        yield AgentResponse(
            source=self.name,
            context=ctx,
            usage=Usage(tokens_input=5, tokens_output=5),
            finish_reason="stop",
        )


@pytest.mark.asyncio
class TestRoundRobinOrchestrator:
    async def test_agent_cycling(self):
        a1 = MockAgent("a1", ["r1"])
        a2 = MockAgent("a2", ["r2"])
        a3 = MockAgent("a3", ["r3"])
        term = MaxMessageTermination(100)
        orchestrator = RoundRobinOrchestrator([a1, a2, a3], term, max_iterations=6)

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

        assert len(orchestrator.shared_messages) == 7
        assert [message.source for message in orchestrator.shared_messages[1:]] == [
            "a1",
            "a2",
            "a3",
            "a1",
            "a2",
            "a3",
        ]
        assert response.pattern_metadata["cycles_completed"] == 2
        assert response.pattern_metadata["agents_order"] == ["a1", "a2", "a3"]

    async def test_default_context_preparation_returns_text_message(self):
        a1 = MockAgent("a1", ["r1"])
        term = MaxMessageTermination(2)
        orchestrator = RoundRobinOrchestrator([a1], term)

        context0 = await orchestrator.prepare_context_for_agent(a1)
        assert isinstance(context0, UserMessage)
        assert "It is now your turn" in context0.content

        orchestrator.shared_messages.append(UserMessage(content="hello", source="user"))
        orchestrator.shared_messages.append(AssistantMessage(content="hi", source="a1"))

        context1 = await orchestrator.prepare_context_for_agent(a1)
        assert isinstance(context1, UserMessage)
        assert "progress/history so far" in context1.content
        assert "hello" in context1.content
        assert "hi" in context1.content

    async def test_recent_message_policy_returns_only_last_n_messages(self):
        a1 = MockAgent("a1", ["r1"])
        term = MaxMessageTermination(2)
        orchestrator = RoundRobinOrchestrator(
            [a1],
            term,
            history_policy=HistoryContextPolicy(
                default=AgentHistoryPolicy(
                    window="recent",
                    message_limit=2,
                    format="messages",
                )
            ),
        )
        orchestrator.shared_messages = [
            UserMessage(content="one", source="user"),
            AssistantMessage(content="two", source="a1"),
            AssistantMessage(content="three", source="a1"),
        ]

        context = await orchestrator.prepare_context_for_agent(a1)

        assert isinstance(context, list)
        assert [message.content for message in context] == ["two", "three"]

    async def test_markdown_and_json_context_formats_are_rendered(self):
        a1 = MockAgent("a1", ["r1"])
        term = MaxMessageTermination(2)
        orchestrator = RoundRobinOrchestrator(
            [a1],
            term,
            history_policy=HistoryContextPolicy(
                default=AgentHistoryPolicy(
                    window="full",
                    format="markdown",
                ),
                per_agent={
                    "a1": AgentHistoryPolicy(
                        window="full",
                        format="json",
                    )
                },
            ),
        )
        orchestrator.shared_messages = [
            UserMessage(content="hello", source="user"),
            AssistantMessage(content="hi", source="a1"),
        ]

        context = await orchestrator.prepare_context_for_agent(a1)

        assert isinstance(context, UserMessage)
        assert '"team_context"' in context.content
        assert '"content": "hello"' in context.content

    async def test_state_update_skips_echo_for_text_context(self):
        a1 = MockAgent("a1", ["r1"])
        term = MaxMessageTermination(2)
        orchestrator = RoundRobinOrchestrator([a1], term)
        orchestrator._last_context_was_text = True

        ctx = AgentContext()
        msg_echo = UserMessage(content="context", source="orchestrator")
        msg_resp = AssistantMessage(content="response", source="a1")
        ctx.messages = [msg_echo, msg_resp]

        result = AgentResponse(source="a1", context=ctx, finish_reason="stop")

        await orchestrator.update_shared_state(result)

        assert len(orchestrator.shared_messages) == 1
        assert orchestrator.shared_messages[0].content == "response"

    async def test_state_update_skips_sent_message_history(self):
        a1 = MockAgent("a1", ["r1"])
        term = MaxMessageTermination(2)
        orchestrator = RoundRobinOrchestrator(
            [a1],
            term,
            history_policy=HistoryContextPolicy(
                default=AgentHistoryPolicy(
                    window="recent",
                    message_limit=2,
                    format="messages",
                )
            ),
        )
        orchestrator._last_context_size = 2
        orchestrator._last_context_was_text = False

        ctx = AgentContext()
        ctx.messages = [
            UserMessage(content="user-1", source="user"),
            AssistantMessage(content="assistant-1", source="a1"),
            AssistantMessage(content="response", source="a1"),
        ]

        result = AgentResponse(source="a1", context=ctx, finish_reason="stop")

        await orchestrator.update_shared_state(result)

        assert [message.content for message in orchestrator.shared_messages] == ["response"]

    async def test_reset(self):
        a1 = MockAgent("a1", ["r1"])
        a2 = MockAgent("a2", ["r2"])
        term = MaxMessageTermination(10)
        orchestrator = RoundRobinOrchestrator([a1, a2], term)

        await orchestrator.select_next_agent()
        assert orchestrator.current_agent_index == 1

        orchestrator._reset_for_run()
        assert orchestrator.current_agent_index == 0
