import pytest
import asyncio
from typing import List, Union, AsyncGenerator, Any

from agentbyte.orchestration.base import BaseOrchestrator
from agentbyte.orchestration.policies import FinalResultPolicy
from agentbyte.agents.base import BaseAgent
from agentbyte.agents.types import AgentResponse, Usage
from agentbyte.messages import Message, UserMessage, AssistantMessage
from agentbyte.termination.max_message import MaxMessageTermination
from agentbyte.types import OrchestrationResponse, OrchestrationEvent
from agentbyte.context import AgentContext


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, context=None, cancellation_token=None) -> AgentResponse:
        resp_text = self.responses[self._resp_idx % len(self.responses)]
        self._resp_idx += 1

        ctx = context or AgentContext()
        msg = AssistantMessage(content=resp_text, source=self.name)
        ctx.messages.append(msg)

        return AgentResponse(
            source=self.name,
            context=ctx,
            usage=Usage(tokens_input=10, tokens_output=10),
            finish_reason="stop",
        )

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

        ctx = context or AgentContext()
        # First yield the user message (context echo) if we want to mimic real agent,
        # but BaseOrchestrator filters them anyway.
        # Let's just yield the assistant 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=10, tokens_output=10),
            finish_reason="stop",
        )


class SimpleOrchestrator(BaseOrchestrator):
    async def select_next_agent(self) -> BaseAgent:
        return self.agents[self.iteration_count % len(self.agents)]

    async def prepare_context_for_agent(self, agent: BaseAgent) -> str:
        return "History: " + " ".join(
            [m.content for m in self.shared_messages if hasattr(m, "content")]
        )

    async def update_shared_state(self, result: AgentResponse) -> None:
        self.shared_messages.append(result.messages[-1])


@pytest.mark.asyncio
class TestBaseOrchestrator:
    async def test_validation(self):
        agent = MockAgent("agent1", ["resp"])
        term = MaxMessageTermination(5)

        # OK
        SimpleOrchestrator([agent], term)

        # Empty agents
        with pytest.raises(ValueError, match="At least one agent is required"):
            SimpleOrchestrator([], term)

        # Duplicate names
        with pytest.raises(ValueError, match="Agent names must be unique"):
            SimpleOrchestrator([agent, MockAgent("agent1", ["other"])], term)

    async def test_basic_run(self):
        agent1 = MockAgent("a1", ["hello"])
        agent2 = MockAgent("a2", ["world"])
        term = MaxMessageTermination(3)  # Initial + 2 more

        orchestrator = SimpleOrchestrator([agent1, agent2], term)
        task = UserMessage(content="Start", source="user")

        response = await orchestrator.run(task)

        assert isinstance(response, OrchestrationResponse)
        assert len(response.messages) == 3  # User + a1 + a2
        assert response.messages[1].content == "hello"
        assert response.messages[2].content == "world"
        assert response.usage.tokens_input == 20
        assert response.usage.tokens_output == 20
        assert "Maximum messages reached" in response.stop_message.content

    async def test_run_stream(self):
        agent1 = MockAgent("a1", ["hello"])
        term = MaxMessageTermination(2)
        orchestrator = SimpleOrchestrator([agent1], term)

        items = []
        async for item in orchestrator.run_stream(
            UserMessage(content="task", source="user")
        ):
            items.append(item)

        # Items should be: UserMessage, AssistantMessage, OrchestrationResponse
        assert len(items) == 3
        assert isinstance(items[0], UserMessage)
        assert isinstance(items[1], AssistantMessage)
        assert isinstance(items[2], OrchestrationResponse)

    async def test_run_accepts_string_task(self):
        agent1 = MockAgent("a1", ["hello"])
        term = MaxMessageTermination(2)
        orchestrator = SimpleOrchestrator([agent1], term)

        response = await orchestrator.run("task")

        assert isinstance(response, OrchestrationResponse)
        assert isinstance(response.messages[0], UserMessage)
        assert response.messages[0].source == "user"
        assert response.messages[0].content == "task"

    async def test_verbose_mode(self):
        agent1 = MockAgent("a1", ["hello"])
        term = MaxMessageTermination(2)
        orchestrator = SimpleOrchestrator([agent1], term)

        events = []
        async for item in orchestrator.run_stream(
            UserMessage(content="task", source="user"), verbose=True
        ):
            if isinstance(item, OrchestrationEvent):
                events.append(item)

        # Expected events: Start, Selection, ExecutionStart, ExecutionComplete, Complete
        event_types = [type(e).__name__ for e in events]
        assert "OrchestrationStartEvent" in event_types
        assert "AgentSelectionEvent" in event_types
        assert "AgentExecutionStartEvent" in event_types
        assert "AgentExecutionCompleteEvent" in event_types
        assert "OrchestrationCompleteEvent" in event_types

    async def test_max_iterations(self):
        agent1 = MockAgent("a1", ["hello"])
        term = MaxMessageTermination(100)  # Won't hit this
        orchestrator = SimpleOrchestrator([agent1], term, max_iterations=2)

        response = await orchestrator.run(UserMessage(content="task", source="user"))
        assert orchestrator.iteration_count == 2
        assert "Maximum iterations reached" in response.stop_message.content

    async def test_cancellation(self):
        from agentbyte.cancellation_token import CancellationToken

        agent1 = MockAgent("a1", ["hello"])
        term = MaxMessageTermination(10)
        orchestrator = SimpleOrchestrator([agent1], term)

        token = CancellationToken()
        token.cancel()

        # Should raise CancelledError or return partial result depending on where it's checked
        # BaseOrchestrator.run catches it and re-raises, but run_stream yields partial result before re-raising

        with pytest.raises(asyncio.CancelledError):
            await orchestrator.run(
                UserMessage(content="task", source="user"), cancellation_token=token
            )

        # Verify it stopped early
        assert orchestrator.iteration_count == 0


def _make_orchestrator(policy: FinalResultPolicy | None = None) -> SimpleOrchestrator:
    agent = MockAgent("a1", ["resp"])
    term = MaxMessageTermination(5)
    orchestrator = SimpleOrchestrator([agent], term)
    orchestrator.final_result_policy = policy
    return orchestrator


class TestFinalResultPolicy:
    def test_no_policy_returns_last_assistant_message(self) -> None:
        orch = _make_orchestrator()
        orch.shared_messages = [
            AssistantMessage(content="first", source="researcher"),
            AssistantMessage(content="last", source="writer"),
        ]
        assert orch._generate_final_result() == "last"

    def test_prefer_agent_selects_correct_agent(self) -> None:
        orch = _make_orchestrator(FinalResultPolicy(prefer_agent="writer"))
        orch.shared_messages = [
            AssistantMessage(content="research notes", source="researcher"),
            AssistantMessage(content="polished article", source="writer"),
            AssistantMessage(content="APPROVED", source="reviewer"),
        ]
        assert orch._generate_final_result() == "polished article"

    def test_prefer_agent_falls_back_when_no_match(self) -> None:
        orch = _make_orchestrator(FinalResultPolicy(prefer_agent="writer"))
        orch.shared_messages = [
            AssistantMessage(content="only researcher", source="researcher"),
        ]
        assert orch._generate_final_result() == "only researcher"

    def test_strip_signals_trailing_removes_last_line(self) -> None:
        orch = _make_orchestrator(
            FinalResultPolicy(prefer_agent="writer", strip_signals=["TERMINATE"])
        )
        orch.shared_messages = [
            AssistantMessage(
                content="Great article about solar energy.\n\nTERMINATE", source="writer"
            ),
        ]
        assert orch._generate_final_result() == "Great article about solar energy."

    def test_strip_signals_trailing_case_insensitive(self) -> None:
        orch = _make_orchestrator(
            FinalResultPolicy(prefer_agent="writer", strip_signals=["TERMINATE"])
        )
        orch.shared_messages = [
            AssistantMessage(content="Draft done.\n\nterminate", source="writer"),
        ]
        assert orch._generate_final_result() == "Draft done."

    def test_strip_signals_trailing_does_not_strip_mid_content(self) -> None:
        orch = _make_orchestrator(
            FinalResultPolicy(prefer_agent="writer", strip_signals=["TERMINATE"])
        )
        orch.shared_messages = [
            AssistantMessage(
                content="TERMINATE is a word we discussed.\n\nConclusion here.",
                source="writer",
            ),
        ]
        result = orch._generate_final_result()
        assert "TERMINATE" in result
        assert "Conclusion here." in result

    def test_strip_signals_anywhere_removes_all_occurrences(self) -> None:
        orch = _make_orchestrator(
            FinalResultPolicy(
                prefer_agent="writer",
                strip_signals=["APPROVED"],
                strip_mode="anywhere",
            )
        )
        orch.shared_messages = [
            AssistantMessage(
                content="The draft was APPROVED. Final: APPROVED done.", source="writer"
            ),
        ]
        result = orch._generate_final_result()
        assert "APPROVED" not in result

    def test_strip_signals_anywhere_is_case_insensitive(self) -> None:
        orch = _make_orchestrator(
            FinalResultPolicy(
                prefer_agent="writer",
                strip_signals=["APPROVED"],
                strip_mode="anywhere",
            )
        )
        orch.shared_messages = [
            AssistantMessage(
                content="The draft was Approved. Also aPpRoVeD.",
                source="writer",
            ),
        ]
        result = orch._generate_final_result()
        assert "Approved" not in result
        assert "aPpRoVeD" not in result

    def test_no_messages_returns_sentinel(self) -> None:
        orch = _make_orchestrator(FinalResultPolicy(prefer_agent="writer"))
        assert orch._generate_final_result() == "No messages generated"

    def test_no_assistant_messages_returns_sentinel(self) -> None:
        orch = _make_orchestrator()
        orch.shared_messages = [
            UserMessage(content="task", source="user"),
        ]
        assert orch._generate_final_result() == "Task completed"

    def test_prefer_agent_case_insensitive(self) -> None:
        orch = _make_orchestrator(FinalResultPolicy(prefer_agent="Writer"))
        orch.shared_messages = [
            AssistantMessage(content="result", source="writer"),
        ]
        assert orch._generate_final_result() == "result"
