import pytest

from logos_adk.providers.mock import MockProvider
from logos_adk.agent import Agent
from logos_adk.crew import CrewTask
from logos_adk.runnable import Runnable
from logos_adk.runtime import Supervisor
from logos_adk.team import Team
from logos_adk.types import Response, Usage


def _make_agent(name: str, response_text: str = "ok") -> Agent:
    provider = MockProvider(responses=[
        Response(
            content=response_text,
            tool_calls=[],
            usage=Usage(input_tokens=1, output_tokens=1),
            raw={},
        ),
    ])
    return Agent(name, model=provider)


class TestTeamCreation:
    def test_team_is_runnable(self):
        team = Team("t", agents=[])
        assert isinstance(team, Runnable)

    def test_team_name(self):
        team = Team("my_team", agents=[])
        assert team.name == "my_team"

    def test_team_agents_registered(self):
        a = _make_agent("a")
        b = _make_agent("b")
        team = Team("t", agents=[a, b])
        assert team.get_agent("a") is a
        assert team.get_agent("b") is b

    def test_duplicate_agent_name_raises(self):
        a1 = _make_agent("dup")
        a2 = _make_agent("dup")
        with pytest.raises(ValueError, match="duplicate"):
            Team("t", agents=[a1, a2])

    def test_orchestrator_must_be_in_agents(self):
        a = _make_agent("a")
        with pytest.raises(ValueError, match="not in agents"):
            Team("t", agents=[a], orchestrator="missing")

    def test_default_orchestrator_is_first_agent(self):
        a = _make_agent("first")
        b = _make_agent("second")
        team = Team("t", agents=[a, b])
        assert team._orchestrator_name == "first"

    def test_explicit_orchestrator(self):
        a = _make_agent("a")
        b = _make_agent("b")
        team = Team("t", agents=[a, b], orchestrator="b")
        assert team._orchestrator_name == "b"

    def test_team_with_supervisor(self):
        sup = Supervisor(on_failure="skip")
        team = Team("t", agents=[], supervisor=sup)
        assert team._config.supervisor is sup

    def test_team_max_rounds(self):
        team = Team("t", agents=[], max_rounds=5)
        assert team._config.max_rounds == 5

    def test_agent_team_backref(self):
        a = _make_agent("a")
        team = Team("t", agents=[a])
        assert a._team is team


class TestTeamRun:
    @pytest.mark.asyncio
    async def test_basic_run_returns_orchestrator_response(self):
        a = _make_agent("orchestrator", "result from orchestrator")
        team = Team("t", agents=[a])
        response = await team.run("do something")
        assert response.content == "result from orchestrator"

    @pytest.mark.asyncio
    async def test_run_passes_prompt_to_orchestrator(self):
        provider = MockProvider(responses=[
            Response(content="echo", tool_calls=[], usage=Usage(1, 1), raw={}),
        ])
        a = Agent("orch", model=provider)
        team = Team("t", agents=[a])
        await team.run("my prompt")
        assert len(provider.call_history) == 1
        messages = provider.call_history[0]["messages"]
        user_msgs = [m for m in messages if m.role == "user"]
        assert any("my prompt" in m.content for m in user_msgs)

    @pytest.mark.asyncio
    async def test_empty_team_raises(self):
        team = Team("t", agents=[])
        with pytest.raises(ValueError, match="no agents"):
            await team.run("hi")

    @pytest.mark.asyncio
    async def test_team_stream_events_exposes_single_top_level_run_boundary(self):
        agent = _make_agent("orchestrator", "team result")
        team = Team("t", agents=[agent])

        events = [event async for event in team.stream_events("hello", session_id="s1")]

        assert events[0]["event"] == "graph_started"
        assert events[-1]["event"] == "graph_completed"
        assert sum(1 for event in events if event["event"] == "graph_started") == 1
        assert sum(1 for event in events if event["event"] == "graph_completed") == 1
        assert any(event["event"] == "node_started" for event in events)
        assert any(event["event"] == "node_completed" for event in events)
        assert events[-1]["output"]["content"] == "team result"

    @pytest.mark.asyncio
    async def test_team_shared_budget_routes_later_tasks_to_cheaper_model(self):
        strong = MockProvider(
            responses=[
                Response(content="strong-one", tool_calls=[], usage=Usage(220, 120), raw={}),
            ]
        )
        strong.model_name = "claude-sonnet-4-20250514"
        cheap = MockProvider(
            responses=[
                Response(content="cheap-two", tool_calls=[], usage=Usage(20, 10), raw={}),
            ]
        )
        cheap.model_name = "gpt-4o-mini"
        router = (
            Agent("router", model=strong)
            .fallback(cheap)
            .cost_aware_routing(enabled=True, budget_preference="balanced")
        )
        team = Team("t", agents=[router]).budget(team_budget_usd=0.001)

        result = await team.run_tasks(
            [
                CrewTask(name="one", description="Analyze migration risks", agent_name="router"),
                CrewTask(name="two", description="Analyze migration rollback", agent_name="router"),
            ],
            kickoff="Work the plan",
        )

        assert [item.output for item in result.task_results] == ["strong-one", "cheap-two"]
        assert result.task_results[1].metadata["raw"]["budget"]["team"]["near_limit"] is True
        assert "team_budget" in result.task_results[1].metadata["raw"]["routing"]["reason"]


class TestTeamRounds:
    @pytest.mark.asyncio
    async def test_max_rounds_zero_raises(self):
        """max_rounds=0 means no rounds allowed, so any pending message triggers error."""
        from logos_adk.errors import MaxRoundsExceeded

        a = Agent("a", model=MockProvider(responses=[
            Response(content="r", tool_calls=[], usage=Usage(1, 1), raw={})
        ]))
        team = Team("t", agents=[a], max_rounds=0)
        await team._mailboxes["a"].put("msg")
        with pytest.raises(MaxRoundsExceeded):
            await team.run_rounds()

    @pytest.mark.asyncio
    async def test_max_rounds_1_completes_if_no_requeue(self):
        """With max_rounds=1, a single round of processing completes normally."""
        a = Agent("a", model=MockProvider(responses=[
            Response(content="done", tool_calls=[], usage=Usage(1, 1), raw={})
        ]))
        team = Team("t", agents=[a], max_rounds=1)
        await team._mailboxes["a"].put("do it")
        response = await team.run_rounds()
        assert response.content == "done"

    @pytest.mark.asyncio
    async def test_rounds_complete_when_mailboxes_empty(self):
        a = _make_agent("a", "done")
        team = Team("t", agents=[a], max_rounds=10)
        response = await team.run("hello")
        assert response.content == "done"
