"""Tests for department-level dynamic orchestration."""
from __future__ import annotations

import pytest

from logos_adk import (
    Agent,
    CrewProcess,
    CrewTask,
    DepartmentSupervisor,
    Graph,
    InMemoryEventBus,
    Swarm,
    Team,
)
from logos_adk.providers.mock import MockProvider
from logos_adk.types import Response, Usage


pytestmark = pytest.mark.asyncio


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


async def _swarm_handler(event):
    return {"content": f"swarm::{event.payload['prompt'].upper()}"}


async def test_department_supervisor_routes_to_selected_team():
    router = Agent(
        "router",
        model=MockProvider(
            responses=[
                (
                    '{"department_name":"seo_department",'
                    '"delegated_prompt":"Optimize the listing title and meta description.",'
                    '"reason":"SEO work should go to the SEO department."}'
                )
            ]
        ),
    )
    seo_provider = _provider("seo result")
    finance_provider = _provider("finance result")
    seo_team = Team("seo-team", agents=[Agent("seo", model=seo_provider)])
    finance_team = Team("finance-team", agents=[Agent("finance", model=finance_provider)])
    supervisor = (
        DepartmentSupervisor("ceo", router=router)
        .department(
            "seo_department",
            seo_team,
            description="Handles SEO optimization, titles, metadata, and ranking issues.",
            capability_tags=["seo", "ranking", "listing"],
        )
        .department(
            "finance_department",
            finance_team,
            description="Handles billing, budgets, and refunds.",
            capability_tags=["finance", "billing"],
        )
    )

    response = await supervisor.run("Need SEO help for this marketplace listing.")

    assert response.content == "seo result"
    assert len(seo_provider.call_history) == 1
    assert len(finance_provider.call_history) == 0
    assert (
        seo_provider.call_history[0]["messages"][-1].content
        == "Optimize the listing title and meta description."
    )
    assert response.raw["dynamic_orchestration"]["department_name"] == "seo_department"
    assert response.raw["dynamic_orchestration"]["department_kind"] == "team"


async def test_department_supervisor_heuristic_fallback_can_choose_crew():
    research_provider = _provider("research done")
    finance_provider = _provider("finance done")
    research_team = Team("research-team", agents=[Agent("researcher", model=research_provider).role("Researcher")])
    research_crew = CrewProcess(
        "research-crew",
        team=research_team,
        tasks=[
            CrewTask(
                name="research",
                description="Investigate SEO and search ranking issues for the product page.",
                role="Researcher",
            )
        ],
    )
    finance_team = Team("finance-team", agents=[Agent("finance", model=finance_provider)])
    supervisor = (
        DepartmentSupervisor("ceo")
        .department(
            "research_crew",
            research_crew,
            description="Researches SEO, ranking changes, search visibility, and marketplace diagnostics.",
            capability_tags=["seo", "research", "ranking"],
        )
        .department(
            "finance_department",
            finance_team,
            description="Handles billing, refunds, invoices, and cost tracking.",
            capability_tags=["finance", "billing", "refund"],
        )
    )

    response = await supervisor.run("Investigate why search ranking dropped for this listing.")

    assert response.content == "research done"
    assert len(research_provider.call_history) == 1
    assert len(finance_provider.call_history) == 0
    assert response.raw["dynamic_orchestration"]["department_name"] == "research_crew"
    assert response.raw["dynamic_orchestration"]["department_kind"] == "crew"
    assert response.raw["dynamic_orchestration"]["reason"] == "Heuristic routing fallback."


async def test_department_supervisor_can_dispatch_to_swarm():
    bus = InMemoryEventBus()
    swarm = Swarm("seo-swarm", event_bus=bus)
    swarm.subscribe("need_seo", _swarm_handler, publish_to="seo_done")
    supervisor = DepartmentSupervisor("ceo").department(
        "seo_swarm",
        swarm,
        description="Parallel SEO swarm for listing optimization and copy tweaks.",
        capability_tags=["seo", "listing", "copy"],
        request_topic="need_seo",
        response_topic="seo_done",
    )

    response = await supervisor.run("optimize seo headline")

    assert response.content == "swarm::OPTIMIZE SEO HEADLINE"
    assert response.raw["dynamic_orchestration"]["department_name"] == "seo_swarm"
    assert response.raw["dynamic_orchestration"]["department_kind"] == "swarm"
    assert len(bus.list_events(topic="need_seo")) == 1
    assert len(bus.list_events(topic="seo_done")) == 1


async def test_department_supervisor_stream_events_serializes_graph_department_events():
    async def draft(ctx):
        return "draft ready"

    graph = Graph("strategy_graph").node(draft, name="draft").terminal("draft")
    supervisor = DepartmentSupervisor("ceo").department(
        "strategy_review",
        graph,
        description="Structured strategy review graph.",
        capability_tags=["strategy", "review"],
    )

    events = [event async for event in supervisor.stream_events("prepare strategy", session_id="sess-1")]

    assert all(isinstance(event, dict) for event in events)
    assert events[0]["event"] == "graph_started"
    assert events[1]["event"] == "node_started"
    assert events[1]["node_name"] == "router"
    assert events[2]["event"] == "node_completed"
    assert events[2]["metadata"]["user_facing"]["step_key"] == "route:select"
    assert any(
        event["event"] == "node_completed"
        and event["node_name"] == "draft"
        and event["metadata"]["department_name"] == "strategy_review"
        for event in events[3:]
    )
    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 events[-1]["event"] == "graph_completed"
    assert events[-1]["metadata"]["department_name"] == "strategy_review"


async def test_department_supervisor_marks_missing_child_terminal_as_failed():
    class SilentRunnable:
        name = "silent"

        async def stream_ui_events(self, prompt: str, **kwargs):
            yield {
                "event": "node_started",
                "event_type": "node",
                "node_name": "silent_step",
                "status": "running",
            }

    supervisor = DepartmentSupervisor("ceo").department(
        "silent_department",
        SilentRunnable(),
        description="Broken routed department without terminal event.",
        capability_tags=["broken"],
    )

    events = [event async for event in supervisor.stream_events("broken flow", session_id="sess-2")]

    assert events[-1]["event"] == "graph_completed"
    assert events[-1]["status"] == "failed"
    assert events[-1]["metadata"]["department_name"] == "silent_department"
    assert events[-1]["metadata"]["user_facing"]["step_key"] == "run:department_missing_terminal"
    assert events[-1]["output"]["error"] == "department_missing_terminal"
