"""Tests for Crew-style process orchestration over Team."""
from __future__ import annotations

import asyncio
import sqlite3
import time
from types import SimpleNamespace

import pytest

from logos_adk.agent import Agent
from logos_adk.crew_checkpoint_store import SQLiteCrewCheckpointStore
from logos_adk.crew import CrewProcess, CrewTask
from logos_adk.providers.mock import MockProvider
from logos_adk.team import Team
from logos_adk.types import Response, Usage


def _provider(texts: list[str] | str) -> MockProvider:
    responses = texts if isinstance(texts, list) else [texts]
    return MockProvider(
        responses=[
            Response(content=text, tool_calls=[], usage=Usage(1, 1), raw={})
            for text in responses
        ]
    )


class DelayedAgent:
    def __init__(self, name: str, role: str, response_text: str, delay: float) -> None:
        self.name = name
        self._response_text = response_text
        self._delay = delay
        self._config = SimpleNamespace(role=role, goal=None, backstory=None)

    @property
    def crew_profile(self) -> dict[str, str]:
        return {"role": self._config.role}

    async def run(self, prompt: str, *, session_id: str | None = None, **kwargs):  # noqa: ARG002
        await asyncio.sleep(self._delay)
        return Response(
            content=self._response_text,
            tool_calls=[],
            usage=Usage(1, 1),
            raw={"prompt": prompt},
        )


@pytest.mark.asyncio
async def test_team_run_tasks_sequential_uses_role_handoff_and_context():
    researcher = Agent("researcher", model=_provider("facts: launch readiness is green")).role("Researcher")
    writer_provider = _provider("brief: launch readiness is green")
    writer = Agent("writer", model=writer_provider).role("Writer")
    team = Team("crew", agents=[researcher, writer], topology="mesh")

    result = await team.run_tasks(
        [
            CrewTask(
                name="research",
                description="Collect evidence for the launch readiness decision.",
                role="Researcher",
                output_key="research_notes",
                handoff_to_role="Writer",
            ),
            CrewTask(
                name="write",
                description="Write the launch brief.",
                context=["research_notes"],
                expected_output="A short brief for executives.",
            ),
        ],
        kickoff="Prepare the launch summary",
        process="sequential",
    )

    assert result.process == "sequential"
    assert result.final_output == "brief: launch readiness is green"
    assert result.task_results[0].assigned_agent == "researcher"
    assert result.task_results[1].assigned_agent == "writer"
    assert result.context["research_notes"] == "facts: launch readiness is green"
    assert result.handoffs[0]["to_role"] == "Writer"
    writer_prompt = writer_provider.call_history[0]["messages"][-1].content
    assert "facts: launch readiness is green" in writer_prompt


@pytest.mark.asyncio
async def test_team_run_tasks_hierarchical_routes_via_manager():
    manager = Agent(
        "manager",
        model=_provider(
            [
                '{"role":"Researcher","agent_name":null,"instruction":"Investigate the source material."}',
                '{"role":"Writer","agent_name":null,"instruction":"Draft the final summary."}',
            ]
        ),
    ).role("Manager")
    researcher = Agent("researcher", model=_provider("evidence bundle")).role("Researcher")
    writer = Agent("writer", model=_provider("final narrative")).role("Writer")
    team = Team("crew", agents=[manager, researcher, writer], orchestrator="manager", topology="mesh")

    result = await team.run_tasks(
        [
            CrewTask(name="research", description="Investigate the issue", output_key="facts"),
            CrewTask(name="write", description="Write the answer", context=["facts"]),
        ],
        kickoff="Resolve the customer issue",
        process="hierarchical",
        manager_role="Manager",
    )

    assert result.process == "hierarchical"
    assert result.manager_agent == "manager"
    assert result.task_results[0].assigned_role == "Researcher"
    assert result.task_results[1].assigned_role == "Writer"
    assert isinstance(result.task_results[0].metadata["raw"], dict)
    assert result.handoffs[0]["manager_agent"] == "manager"
    assert result.final_output == "final narrative"


@pytest.mark.asyncio
async def test_crew_process_wraps_team_tasks_as_runnable():
    researcher = Agent("researcher", model=_provider("research done")).role("Researcher")
    team = Team("crew", agents=[researcher], topology="mesh")
    process = CrewProcess(
        "crew-process",
        team=team,
        tasks=[CrewTask(name="research", description="Research the topic", role="Researcher")],
    )

    response = await process.run("Investigate the topic")

    assert response.content == "research done"
    assert response.raw["crew_result"]["process"] == "sequential"
    assert response.raw["crew_result"]["task_results"][0]["assigned_agent"] == "researcher"


@pytest.mark.asyncio
async def test_team_run_tasks_executes_parallel_ready_tasks_with_dependencies():
    agent_a = DelayedAgent("research_a", "ResearcherA", "facts-a", 0.06)
    agent_b = DelayedAgent("research_b", "ResearcherB", "facts-b", 0.06)
    writer_provider = _provider("merged brief")
    writer = Agent("writer", model=writer_provider).role("Writer")
    team = Team("crew", agents=[agent_a, agent_b, writer], topology="mesh")

    started = time.monotonic()
    result = await team.run_tasks(
        [
            CrewTask(
                name="research_a",
                description="Collect source A",
                agent_name="research_a",
                async_execution=True,
                output_key="facts_a",
            ),
            CrewTask(
                name="research_b",
                description="Collect source B",
                agent_name="research_b",
                async_execution=True,
                output_key="facts_b",
            ),
            CrewTask(
                name="write",
                description="Write the combined brief",
                role="Writer",
                depends_on=["research_a", "research_b"],
                context=["facts_a", "facts_b"],
            ),
        ],
        kickoff="Prepare the package",
        process="sequential",
        max_parallel_tasks=2,
    )
    elapsed = time.monotonic() - started

    assert result.status == "completed"
    assert result.final_output == "merged brief"
    assert elapsed < 0.11
    writer_prompt = writer_provider.call_history[0]["messages"][-1].content
    assert "facts-a" in writer_prompt
    assert "facts-b" in writer_prompt


@pytest.mark.asyncio
async def test_team_run_tasks_pauses_on_approval_task_without_decision():
    researcher = Agent("researcher", model=_provider("facts gathered")).role("Researcher")
    writer = Agent("writer", model=_provider("final draft")).role("Writer")
    team = Team("crew", agents=[researcher, writer], topology="mesh")

    result = await team.run_tasks(
        [
            CrewTask(name="research", description="Gather facts", role="Researcher", output_key="facts"),
            CrewTask(
                name="review",
                description="Approve the gathered facts",
                task_type="approval",
                depends_on=["research"],
                approval_message="Human review required",
            ),
            CrewTask(
                name="write",
                description="Write the final draft",
                role="Writer",
                depends_on=["review"],
                context=["facts"],
            ),
        ],
        kickoff="Prepare the answer",
        process="sequential",
    )

    assert result.status == "paused"
    assert result.run_id
    assert result.pending_approval is not None
    assert result.pending_approval["run_id"] == result.run_id
    assert result.pending_approval["task_name"] == "review"
    assert result.final_output == "facts gathered"
    assert [item.task_name for item in result.task_results] == ["research", "review"]
    assert result.task_results[-1].status == "paused"


@pytest.mark.asyncio
async def test_team_run_tasks_continues_when_approval_is_supplied():
    researcher = Agent("researcher", model=_provider("facts gathered")).role("Researcher")
    writer = Agent("writer", model=_provider("final draft")).role("Writer")
    team = Team("crew", agents=[researcher, writer], topology="mesh")

    result = await team.run_tasks(
        [
            CrewTask(name="research", description="Gather facts", role="Researcher", output_key="facts"),
            CrewTask(
                name="review",
                description="Approve the gathered facts",
                task_type="approval",
                depends_on=["research"],
                approval_message="Human review required",
            ),
            CrewTask(
                name="write",
                description="Write the final draft",
                role="Writer",
                depends_on=["review"],
                context=["facts"],
            ),
        ],
        kickoff="Prepare the answer",
        process="sequential",
        approval_responses={"review": True},
    )

    assert result.status == "completed"
    assert result.context["review"] == "approved"
    assert result.final_output == "final draft"
    assert [item.task_name for item in result.task_results] == ["research", "review", "write"]


@pytest.mark.asyncio
async def test_team_resume_approval_continues_from_paused_checkpoint():
    researcher_provider = _provider("facts gathered")
    researcher = Agent("researcher", model=researcher_provider).role("Researcher")
    writer_provider = _provider("final draft")
    writer = Agent("writer", model=writer_provider).role("Writer")
    team = Team("crew", agents=[researcher, writer], topology="mesh")
    tasks = [
        CrewTask(name="research", description="Gather facts", role="Researcher", output_key="facts"),
        CrewTask(
            name="review",
            description="Approve the gathered facts",
            task_type="approval",
            depends_on=["research"],
            approval_message="Human review required",
        ),
        CrewTask(
            name="write",
            description="Write the final draft",
            role="Writer",
            depends_on=["review"],
            context=["facts"],
        ),
    ]

    paused = await team.run_tasks(tasks, kickoff="Prepare the answer", process="sequential")
    resumed = await team.resume_approval(paused.run_id, task_name="review", approved=True)

    assert paused.status == "paused"
    assert resumed.run_id == paused.run_id
    assert resumed.status == "completed"
    assert resumed.context["review"] == "approved"
    assert resumed.final_output == "final draft"
    assert [item.task_name for item in resumed.task_results] == ["research", "review", "write"]
    assert len(researcher_provider.call_history) == 1
    assert len(writer_provider.call_history) == 1


@pytest.mark.asyncio
async def test_crew_process_resume_approval_uses_last_paused_run():
    researcher = Agent("researcher", model=_provider("facts gathered")).role("Researcher")
    writer = Agent("writer", model=_provider("final draft")).role("Writer")
    team = Team("crew", agents=[researcher, writer], topology="mesh")
    process = CrewProcess(
        "crew-process",
        team=team,
        tasks=[
            CrewTask(name="research", description="Gather facts", role="Researcher", output_key="facts"),
            CrewTask(
                name="review",
                description="Approve the gathered facts",
                task_type="approval",
                depends_on=["research"],
                approval_message="Human review required",
            ),
            CrewTask(
                name="write",
                description="Write the final draft",
                role="Writer",
                depends_on=["review"],
                context=["facts"],
            ),
        ],
    )

    paused = await process.run("Prepare the answer")
    resumed = await process.resume_approval(task_name="review", approved=True)

    assert paused.raw["crew_result"]["status"] == "paused"
    assert resumed.raw["crew_result"]["status"] == "completed"
    assert resumed.raw["crew_result"]["run_id"] == paused.raw["crew_result"]["run_id"]
    assert resumed.content == "final draft"


@pytest.mark.asyncio
async def test_team_resume_approval_survives_restart_with_sqlite_checkpoint_store(tmp_path):
    store = SQLiteCrewCheckpointStore(str(tmp_path / "crew_checkpoints.db"))
    before_research_provider = _provider("facts gathered")
    before_writer_provider = _provider("draft before restart should not run")
    team_before_restart = Team(
        "crew",
        agents=[
            Agent("researcher", model=before_research_provider).role("Researcher"),
            Agent("writer", model=before_writer_provider).role("Writer"),
        ],
        topology="mesh",
        checkpoint_store=store,
    )
    tasks = [
        CrewTask(name="research", description="Gather facts", role="Researcher", output_key="facts"),
        CrewTask(
            name="review",
            description="Approve the gathered facts",
            task_type="approval",
            depends_on=["research"],
            approval_message="Human review required",
        ),
        CrewTask(
            name="write",
            description="Write the final draft",
            role="Writer",
            depends_on=["review"],
            context=["facts"],
        ),
    ]

    paused = await team_before_restart.run_tasks(tasks, kickoff="Prepare the answer", process="sequential")

    after_research_provider = _provider("research should not rerun")
    after_writer_provider = _provider("final draft after restart")
    team_after_restart = Team(
        "crew",
        agents=[
            Agent("researcher", model=after_research_provider).role("Researcher"),
            Agent("writer", model=after_writer_provider).role("Writer"),
        ],
        topology="mesh",
        checkpoint_store=SQLiteCrewCheckpointStore(str(tmp_path / "crew_checkpoints.db")),
    )

    resumed = await team_after_restart.resume_approval(paused.run_id, task_name="review", approved=True)

    assert paused.status == "paused"
    assert resumed.status == "completed"
    assert resumed.run_id == paused.run_id
    assert resumed.final_output == "final draft after restart"
    assert resumed.context["review"] == "approved"
    assert len(before_research_provider.call_history) == 1
    assert len(before_writer_provider.call_history) == 0
    assert len(after_research_provider.call_history) == 0
    assert len(after_writer_provider.call_history) == 1


@pytest.mark.asyncio
async def test_team_checkpoint_store_rejects_non_json_task_metadata(tmp_path):
    store = SQLiteCrewCheckpointStore(str(tmp_path / "crew_checkpoints.db"))
    reviewer = CrewTask(
        name="review",
        description="Approve the gathered facts",
        task_type="approval",
        approval_message="Human review required",
        metadata={"bad": object()},
    )
    team = Team("crew", agents=[], topology="mesh", checkpoint_store=store)

    with pytest.raises(ValueError, match=r"metadata\.bad"):
        await team.run_tasks([reviewer], kickoff="Prepare the answer", process="sequential")

    assert store.load_checkpoint("missing-run-id") is None
    with sqlite3.connect(store.path) as connection:
        rows = connection.execute(
            f"SELECT COUNT(*) FROM {store.table_name}"
        ).fetchone()[0]
    assert rows == 0


@pytest.mark.asyncio
async def test_team_checkpoint_store_rejects_non_json_run_kwargs(tmp_path):
    store = SQLiteCrewCheckpointStore(str(tmp_path / "crew_checkpoints.db"))
    reviewer = CrewTask(
        name="review",
        description="Approve the gathered facts",
        task_type="approval",
        approval_message="Human review required",
    )
    team = Team("crew", agents=[], topology="mesh", checkpoint_store=store)

    with pytest.raises(ValueError, match=r"kwargs\.bad"):
        await team.run_tasks(
            [reviewer],
            kickoff="Prepare the answer",
            process="sequential",
            bad=object(),
        )

    assert store.load_checkpoint("missing-run-id") is None
    with sqlite3.connect(store.path) as connection:
        rows = connection.execute(
            f"SELECT COUNT(*) FROM {store.table_name}"
        ).fetchone()[0]
    assert rows == 0
