import asyncio
from dataclasses import dataclass

import pytest

from logos_adk.a2a import A2ARequest
from logos_adk.agent import Agent
from logos_adk.event_bus import InMemoryEventBus
from logos_adk.errors import RoutingError
from logos_adk.handler import handler
from logos_adk.providers.mock import MockProvider
from logos_adk.team import Team
from logos_adk.types import Response, Usage


def _mock_provider(text: str = "ok") -> MockProvider:
    return MockProvider(responses=[
        Response(content=text, tool_calls=[], usage=Usage(1, 1), raw={}),
    ])


class TestTeamSend:
    @pytest.mark.asyncio
    async def test_send_delivers_to_mailbox(self):
        a = Agent("a", model=_mock_provider())
        b = Agent("b", model=_mock_provider())
        team = Team("t", agents=[a, b], topology="mesh")
        await team.send("a", "b", "hello")
        msg = team._mailboxes["b"].get_nowait()
        assert msg == "hello"

    @pytest.mark.asyncio
    async def test_send_blocked_by_topology(self):
        a = Agent("a", model=_mock_provider())
        b = Agent("b", model=_mock_provider())
        team = Team("t", agents=[a, b], topology={"a": ["b"], "b": []})
        with pytest.raises(RoutingError):
            await team.send("b", "a", "nope")

    @pytest.mark.asyncio
    async def test_broadcast_sends_to_all_allowed(self):
        a = Agent("a", model=_mock_provider())
        b = Agent("b", model=_mock_provider())
        c = Agent("c", model=_mock_provider())
        team = Team("t", agents=[a, b, c], topology="mesh")
        await team.broadcast("a", "hi all")
        assert team._mailboxes["b"].get_nowait() == "hi all"
        assert team._mailboxes["c"].get_nowait() == "hi all"
        assert team._mailboxes["a"].get_nowait() is None  # not to self

    @pytest.mark.asyncio
    async def test_broadcast_respects_topology(self):
        a = Agent("a", model=_mock_provider())
        b = Agent("b", model=_mock_provider())
        c = Agent("c", model=_mock_provider())
        team = Team("t", agents=[a, b, c], topology={"a": ["b"], "b": [], "c": []})
        await team.broadcast("a", "msg")
        assert team._mailboxes["b"].get_nowait() == "msg"
        assert team._mailboxes["c"].get_nowait() is None


class TestAgentSendBroadcast:
    @pytest.mark.asyncio
    async def test_agent_send_delegates_to_team(self):
        a = Agent("a", model=_mock_provider())
        b = Agent("b", model=_mock_provider())
        Team("t", agents=[a, b], topology="mesh")
        await a.send_to("b", "message")
        # Access mailbox through Team's internal structure
        assert a._team._mailboxes["b"].get_nowait() == "message"

    @pytest.mark.asyncio
    async def test_agent_send_without_team_raises(self):
        a = Agent("a", model=_mock_provider())
        with pytest.raises(RuntimeError, match="not part of a team"):
            await a.send_to("b", "msg")

    @pytest.mark.asyncio
    async def test_agent_broadcast_delegates_to_team(self):
        a = Agent("a", model=_mock_provider())
        b = Agent("b", model=_mock_provider())
        Team("t", agents=[a, b], topology="mesh")
        await a.broadcast_msg("hi")
        assert a._team._mailboxes["b"].get_nowait() == "hi"

    def test_agent_publish_event_delegates_to_team(self):
        a = Agent("a", model=_mock_provider())
        bus = InMemoryEventBus()
        Team(
            "t",
            agents=[a],
            event_bus=bus,
            topic_policy={"a": ["need_seo"]},
        )

        published = a.publish_event("need_seo", {"prompt": "optimize"})

        assert published.topic == "need_seo"
        events = bus.list_events(topic="need_seo")
        assert len(events) == 1
        assert events[0].publisher == "a"

    def test_agent_schedule_event_delegates_to_team(self):
        a = Agent("a", model=_mock_provider())
        bus = InMemoryEventBus()
        Team(
            "t",
            agents=[a],
            event_bus=bus,
            topic_policy={"a": ["need_seo"]},
        )

        scheduled = a.schedule_event("need_seo", {"prompt": "optimize"}, campaign_id="cmp-1", delay_seconds=0.1)

        assert scheduled.topic == "need_seo"
        events = bus.list_events(topic="need_seo", campaign_id="cmp-1")
        assert len(events) == 1
        assert events[0].campaign_id == "cmp-1"

    def test_agent_schedule_task_delegates_to_team(self):
        a = Agent("a", model=_mock_provider())
        bus = InMemoryEventBus()
        Team("t", agents=[a], event_bus=bus)

        scheduled = a.schedule_task(
            "check ctr tomorrow",
            delay_seconds=0.1,
            campaign_id="cmp-task",
            session_id="session-task",
        )

        assert scheduled.topic == "__scheduled_task__:a"
        events = bus.list_events(topic="__scheduled_task__:a", campaign_id="cmp-task")
        assert len(events) == 1
        assert events[0].payload["task"] == "check ctr tomorrow"
        assert events[0].payload["source_agent"] == "a"
        assert events[0].payload["target_agent"] == "a"
        assert events[0].payload["session_id"] == "session-task"

    def test_agent_publish_event_without_team_raises(self):
        a = Agent("a", model=_mock_provider())
        with pytest.raises(RuntimeError, match="not part of a team"):
            a.publish_event("need_seo", {"prompt": "optimize"})

    def test_agent_schedule_task_without_team_raises(self):
        a = Agent("a", model=_mock_provider())
        with pytest.raises(RuntimeError, match="not part of a team"):
            a.schedule_task("check later")

    @pytest.mark.asyncio
    async def test_agent_request_a2a_returns_protocol_response(self):
        requester = Agent("requester", model=_mock_provider())
        worker = Agent("worker", model=_mock_provider("completed by worker"))
        Team("t", agents=[requester, worker], topology="mesh")

        response = await requester.request_a2a("worker", "Do the task", metadata={"priority": "high"})

        assert response.status == "completed"
        assert response.from_agent == "worker"
        assert response.to_agent == "requester"
        assert response.content == "completed by worker"
        assert response.metadata["task_id"]

    @pytest.mark.asyncio
    async def test_temporal_autonomy_tool_schedules_task(self):
        analyst = Agent("analyst", model=_mock_provider()).temporal_autonomy()
        bus = InMemoryEventBus()
        Team("growth", agents=[analyst], event_bus=bus)

        tool = next(item for item in analyst._config.tools if item.name == "schedule_task")
        result = await tool(
            task="check ctr in 24h",
            delay_seconds=3600,
            campaign_id="cmp-tool",
        )

        assert result["topic"] == "__scheduled_task__:analyst"
        assert result["campaign_id"] == "cmp-tool"
        events = bus.list_events(topic="__scheduled_task__:analyst", campaign_id="cmp-tool")
        assert len(events) == 1
        assert events[0].payload["task"] == "check ctr in 24h"


@dataclass
class ResearchTask:
    query: str


@dataclass
class WriteTask:
    data: str


class TestHandlerDispatch:
    @pytest.mark.asyncio
    async def test_handler_dispatches_typed_message(self):
        """@handler methods are invoked when mailbox receives matching type."""
        received: list[ResearchTask] = []

        class Researcher(Agent):
            @handler(ResearchTask)
            async def on_research(self, task: ResearchTask) -> None:
                received.append(task)

        r = Researcher("researcher", model=_mock_provider())
        w = Agent("writer", model=_mock_provider())
        team = Team("t", agents=[r, w], topology="mesh")

        task = ResearchTask(query="AI trends")
        await team._mailboxes["researcher"].put(task)
        await team.run_rounds()

        assert len(received) == 1
        assert received[0].query == "AI trends"

    @pytest.mark.asyncio
    async def test_handler_not_matched_falls_back_to_run(self):
        """String messages without handler fall back to agent.run()."""
        provider = MockProvider(responses=[
            Response(content="fallback", tool_calls=[], usage=Usage(1, 1), raw={}),
        ])
        a = Agent("a", model=provider)
        team = Team("t", agents=[a])
        await team._mailboxes["a"].put("plain string")
        response = await team.run_rounds()
        assert response.content == "fallback"
        assert len(provider.call_history) == 1

    @pytest.mark.asyncio
    async def test_handler_chain_between_agents(self):
        """Agent A handles a message and sends to Agent B via handler."""
        write_received: list[WriteTask] = []

        class Researcher(Agent):
            @handler(ResearchTask)
            async def on_research(self, task: ResearchTask) -> None:
                await self.send_to("writer", WriteTask(data=f"found: {task.query}"))

        class Writer(Agent):
            @handler(WriteTask)
            async def on_write(self, task: WriteTask) -> None:
                write_received.append(task)

        r = Researcher("researcher", model=_mock_provider())
        w = Writer("writer", model=_mock_provider())
        team = Team("t", agents=[r, w], topology="mesh")

        await team._mailboxes["researcher"].put(ResearchTask(query="quantum"))
        await team.run_rounds()

        assert len(write_received) == 1
        assert write_received[0].data == "found: quantum"

    @pytest.mark.asyncio
    async def test_a2a_handler_can_respond_explicitly(self):
        class Researcher(Agent):
            @handler(A2ARequest)
            async def on_request(self, task: A2ARequest) -> None:
                await self.respond_a2a(task, f"researched: {task.content}", metadata={"source": "handler"})

        requester = Agent("requester", model=_mock_provider())
        researcher = Researcher("researcher", model=_mock_provider())
        Team("t", agents=[requester, researcher], topology="mesh")

        response = await requester.request_a2a("researcher", "AI adoption")

        assert response.content == "researched: AI adoption"
        assert response.metadata["source"] == "handler"


class TestTeamRoles:
    @pytest.mark.asyncio
    async def test_kickoff_can_route_by_role(self):
        researcher = Agent("researcher", model=_mock_provider("role-routed")).role(
            "Researcher",
            goal="Find facts",
        )
        writer = Agent("writer", model=_mock_provider("writer-output")).role("Writer")
        team = Team("crew", agents=[researcher, writer], topology="mesh")

        response = await team.kickoff("Investigate", role="Researcher")

        assert response.content == "role-routed"
        assert team.get_agent_by_role("Writer") is writer

    @pytest.mark.asyncio
    async def test_delegate_can_target_role(self):
        researcher = Agent("researcher", model=_mock_provider("research done")).role("Researcher")
        writer = Agent("writer", model=_mock_provider("draft ready")).role("Writer")
        team = Team("crew", agents=[researcher, writer], topology="mesh")

        response = await team.delegate("Write the brief", role="Writer")

        assert response.content == "draft ready"


class TestTeamEventBus:
    @pytest.mark.asyncio
    async def test_run_rounds_executes_due_scheduled_task_for_self(self):
        provider = _mock_provider("ctr checked")
        analyst = Agent("analyst", model=provider)
        bus = InMemoryEventBus()
        team = Team("growth", agents=[analyst], event_bus=bus)

        team.schedule_task(
            "analyst",
            "check ctr now",
            campaign_id="cmp-self-task",
            correlation_id="corr-self-task",
        )

        response = await team.run_rounds()

        assert response.content == "ctr checked"
        assert len(provider.call_history) == 1
        messages = provider.call_history[0]["messages"]
        assert messages[-1].content == "check ctr now"
        events = bus.list_events(topic="__scheduled_task__:analyst", campaign_id="cmp-self-task")
        assert len(events) == 1
        assert events[0].status == "completed"

    @pytest.mark.asyncio
    async def test_run_rounds_executes_due_scheduled_task_for_other_agent(self):
        strategist = Agent("strategist", model=_mock_provider("unused"))
        copywriter_provider = _mock_provider("copy updated")
        copywriter = Agent("copywriter", model=copywriter_provider)
        bus = InMemoryEventBus()
        team = Team(
            "growth",
            agents=[strategist, copywriter],
            topology="mesh",
            event_bus=bus,
        )

        team.schedule_task(
            "strategist",
            "rewrite headline if ctr is low",
            to_agent="copywriter",
            campaign_id="cmp-cross-task",
            correlation_id="corr-cross-task",
        )

        response = await team.run_rounds()

        assert response.content == "copy updated"
        assert len(copywriter_provider.call_history) == 1
        messages = copywriter_provider.call_history[0]["messages"]
        assert messages[-1].content == "rewrite headline if ctr is low"
        events = bus.list_events(topic="__scheduled_task__:copywriter", campaign_id="cmp-cross-task")
        assert len(events) == 1
        assert events[0].status == "completed"

    def test_schedule_task_respects_topology(self):
        strategist = Agent("strategist", model=_mock_provider())
        copywriter = Agent("copywriter", model=_mock_provider())
        team = Team(
            "growth",
            agents=[strategist, copywriter],
            topology={"strategist": [], "copywriter": []},
            event_bus=InMemoryEventBus(),
        )

        with pytest.raises(RoutingError, match="not allowed by topology"):
            team.schedule_task(
                "strategist",
                "rewrite headline",
                to_agent="copywriter",
            )

    @pytest.mark.asyncio
    async def test_agent_subscribe_topic_before_team_is_registered_on_attach(self):
        seo = Agent("seo", model=_mock_provider("optimized copy"))
        seo.subscribe_topic("need_seo", publish_to="seo_done")
        bus = InMemoryEventBus()
        team = Team(
            "t",
            agents=[seo],
            event_bus=bus,
            topic_policy={"seo": ["need_seo", "seo_done"]},
        )

        assert seo._team is team
        assert len(team._topic_subscriptions) == 2
        assert any(item.topic == "__scheduled_task__:seo" for item in team._topic_subscriptions)
        assert any(item.topic == "need_seo" for item in team._topic_subscriptions)
        seo.publish_event(
            "need_seo",
            {"prompt": "summer sale"},
            correlation_id="corr-pre",
            run_id="run-pre",
        )

        response = await team.run_rounds()

        assert response.content == "optimized copy"
        events = bus.list_events(topic="seo_done", correlation_id="corr-pre", run_id="run-pre")
        assert len(events) == 1
        assert events[0].payload["content"] == "optimized copy"

    @pytest.mark.asyncio
    async def test_agent_subscribe_topic_after_team_registers_immediately(self):
        reviewer = Agent("reviewer", model=_mock_provider("unused"))
        bus = InMemoryEventBus()
        team = Team(
            "t",
            agents=[reviewer],
            event_bus=bus,
            topic_policy={"reviewer": ["need_review", "review_done"]},
        )

        async def handle_review(event):
            return {"approved": event.payload["score"] >= 0.8}

        reviewer.subscribe_topic("need_review", handler=handle_review, publish_to="review_done")
        assert len(team._topic_subscriptions) == 2
        assert any(item.topic == "__scheduled_task__:reviewer" for item in team._topic_subscriptions)
        assert any(item.topic == "need_review" for item in team._topic_subscriptions)

        reviewer.publish_event(
            "need_review",
            {"score": 0.9},
            correlation_id="corr-post",
            run_id="run-post",
        )

        await team.run_rounds()

        events = bus.list_events(topic="review_done", correlation_id="corr-post", run_id="run-post")
        assert len(events) == 1
        assert events[0].payload == {"approved": True}

    def test_publish_event_respects_topic_policy(self):
        a = Agent("a", model=_mock_provider())
        b = Agent("b", model=_mock_provider())
        bus = InMemoryEventBus()
        team = Team(
            "t",
            agents=[a, b],
            event_bus=bus,
            topic_policy={"a": ["need_seo"], "b": []},
        )

        published = team.publish_event("a", "need_seo", {"prompt": "optimize"})
        assert published.topic == "need_seo"

        with pytest.raises(RoutingError, match="cannot publish topic"):
            team.publish_event("b", "need_seo", {"prompt": "blocked"})

    @pytest.mark.asyncio
    async def test_run_rounds_processes_team_event_subscription_and_publishes_response(self):
        seo = Agent("seo", model=_mock_provider("optimized copy"))
        bus = InMemoryEventBus()
        team = Team(
            "t",
            agents=[seo],
            event_bus=bus,
            topic_policy={"seo": ["need_seo", "seo_done"]},
        )
        team.subscribe("seo", "need_seo", publish_to="seo_done")
        team.publish_event(
            "seo",
            "need_seo",
            {"prompt": "summer sale"},
            correlation_id="corr-team",
            run_id="run-team",
        )

        response = await team.run_rounds()

        assert response.content == "optimized copy"
        request_events = bus.list_events(topic="need_seo", correlation_id="corr-team", run_id="run-team")
        assert len(request_events) == 1
        assert request_events[0].status == "completed"
        response_events = bus.list_events(topic="seo_done", correlation_id="corr-team", run_id="run-team")
        assert len(response_events) == 1
        assert response_events[0].payload["content"] == "optimized copy"

    @pytest.mark.asyncio
    async def test_schedule_event_and_campaign_state(self):
        bus = InMemoryEventBus()

        async def handle_seo(event):
            return {"done": event.payload["prompt"]}

        seo = Agent("seo", model=_mock_provider("unused"))
        team = Team(
            "t",
            agents=[seo],
            event_bus=bus,
            topic_policy={"seo": ["need_seo", "seo_done"]},
        )
        team.subscribe("seo", "need_seo", handler=handle_seo, publish_to="seo_done")

        team.schedule_event(
            "seo",
            "need_seo",
            {"prompt": "summer sale"},
            campaign_id="cmp-summer",
            correlation_id="corr-summer",
            run_id="run-summer",
            delay_seconds=0.05,
        )

        initial_state = team.campaign_state("cmp-summer")
        assert initial_state.total_events == 1
        assert initial_state.queued_events == 1

        await team.run_rounds()
        state_before_due = team.campaign_state("cmp-summer")
        assert state_before_due.completed_events == 0

        await asyncio.sleep(0.06)
        await team.run_rounds()

        state_after_due = team.campaign_state("cmp-summer")
        assert state_after_due.total_events == 2
        assert state_after_due.completed_events == 1
        response_events = bus.list_events(topic="seo_done", campaign_id="cmp-summer")
        assert len(response_events) == 1
        assert response_events[0].campaign_id == "cmp-summer"

    @pytest.mark.asyncio
    async def test_run_rounds_processes_custom_subscription_handler(self):
        reviewer = Agent("reviewer", model=_mock_provider("unused"))
        bus = InMemoryEventBus()
        team = Team(
            "t",
            agents=[reviewer],
            event_bus=bus,
            topic_policy={"reviewer": ["need_review", "review_done"]},
        )

        async def handle_review(event):
            return {"approved": event.payload["score"] >= 0.9}

        team.subscribe("reviewer", "need_review", handler=handle_review, publish_to="review_done")
        team.publish_event(
            "reviewer",
            "need_review",
            {"score": 0.95},
            correlation_id="corr-review",
            run_id="run-review",
        )

        response = await team.run_rounds()

        assert "\"approved\": true" in response.content.lower()
        response_events = bus.list_events(topic="review_done", correlation_id="corr-review", run_id="run-review")
        assert len(response_events) == 1
        assert response_events[0].payload == {"approved": True}
