"""Tests for the Task Manager."""

from __future__ import annotations

import pytest

from agent_scout.db.models import Task, TaskStatus
from agent_scout.events import Event, EventBus, EventType
from agent_scout.task_manager import TaskManager

# --- Mock Agent ---


class MockAgent:
    """Mock agent that completes tasks immediately."""

    def __init__(self, delay: float = 0.0) -> None:
        self._delay = delay
        self.ran_tasks: list[str] = []
        self._active: dict[str, bool] = {}

    async def run_task(self, task: Task) -> Task:
        import asyncio

        self.ran_tasks.append(task.id)
        self._active[task.id] = True
        if self._delay:
            await asyncio.sleep(self._delay)
        self._active.pop(task.id, None)
        task.status = TaskStatus.COMPLETED
        task.result = {"done": True}
        return task

    async def stop_task(self, task_id: str) -> None:
        self._active.pop(task_id, None)

    @property
    def active_task_ids(self) -> list[str]:
        return list(self._active.keys())


class FailingAgent(MockAgent):
    """Mock agent that fails tasks."""

    async def run_task(self, task: Task) -> Task:
        self.ran_tasks.append(task.id)
        task.status = TaskStatus.FAILED
        task.error = "Simulated failure"
        return task


@pytest.fixture
def bus():
    return EventBus()


# --- Tests ---


async def test_submit_task(bus: EventBus):
    agent = MockAgent()
    tm = TaskManager(agent, bus)

    task = await tm.submit("Find flights")
    assert task.status == TaskStatus.PENDING
    assert task.description == "Find flights"
    assert tm.pending_count == 1


async def test_execute_task(bus: EventBus):
    agent = MockAgent()
    tm = TaskManager(agent, bus)

    task = await tm.submit("Find flights")
    result = await tm.execute(task.id)

    assert result.status == TaskStatus.COMPLETED
    assert result.result == {"done": True}
    assert task.id in agent.ran_tasks


async def test_execute_nonexistent_task(bus: EventBus):
    agent = MockAgent()
    tm = TaskManager(agent, bus)

    with pytest.raises(ValueError, match="not found"):
        await tm.execute("nonexistent")


async def test_execute_all(bus: EventBus):
    agent = MockAgent()
    tm = TaskManager(agent, bus)

    await tm.submit("Task 1")
    await tm.submit("Task 2")
    await tm.submit("Task 3")

    results = await tm.execute_all()
    assert len(results) == 3
    assert all(r.status == TaskStatus.COMPLETED for r in results)


async def test_priority_ordering(bus: EventBus):
    agent = MockAgent()
    tm = TaskManager(agent, bus)

    t_low = await tm.submit("Low priority", priority=1)
    t_high = await tm.submit("High priority", priority=10)
    t_mid = await tm.submit("Mid priority", priority=5)

    tasks = tm.list_tasks(status=TaskStatus.PENDING)
    assert tasks[0].id == t_high.id
    assert tasks[1].id == t_mid.id
    assert tasks[2].id == t_low.id


async def test_concurrency_limit(bus: EventBus):
    """Verify max concurrent tasks is respected."""
    import asyncio

    max_concurrent_seen = 0
    current_concurrent = 0
    lock = asyncio.Lock()

    class ConcurrencyTrackingAgent:
        async def run_task(self, task: Task) -> Task:
            nonlocal max_concurrent_seen, current_concurrent
            async with lock:
                current_concurrent += 1
                max_concurrent_seen = max(max_concurrent_seen, current_concurrent)
            await asyncio.sleep(0.02)
            async with lock:
                current_concurrent -= 1
            task.status = TaskStatus.COMPLETED
            task.result = {"done": True}
            return task

        async def stop_task(self, task_id): ...
        @property
        def active_task_ids(self):
            return []

    tm = TaskManager(ConcurrencyTrackingAgent(), bus, max_concurrent=2)

    for i in range(5):
        await tm.submit(f"Task {i}")

    await tm.execute_all()
    assert max_concurrent_seen <= 2


async def test_task_dependencies(bus: EventBus):
    """Tasks with dependencies wait for their dependencies."""
    import asyncio

    execution_order = []

    class OrderTrackingAgent:
        async def run_task(self, task: Task) -> Task:
            execution_order.append(task.id)
            await asyncio.sleep(0.01)
            task.status = TaskStatus.COMPLETED
            task.result = {"done": True}
            return task

        async def stop_task(self, task_id): ...
        @property
        def active_task_ids(self):
            return []

    tm = TaskManager(OrderTrackingAgent(), bus)

    t1 = await tm.submit("First task")
    t2 = await tm.submit("Second task (depends on first)", depends_on=[t1.id])

    # Execute both
    await tm.execute_all()

    # t1 should have completed before t2 started
    assert t1.id in execution_order
    assert t2.id in execution_order
    idx1 = execution_order.index(t1.id)
    idx2 = execution_order.index(t2.id)
    assert idx1 < idx2


async def test_cancel_pending_task(bus: EventBus):
    agent = MockAgent()
    tm = TaskManager(agent, bus)

    task = await tm.submit("Cancellable task")
    await tm.cancel(task.id)

    assert tm.get_task(task.id).status == TaskStatus.CANCELLED


async def test_list_tasks_by_status(bus: EventBus):
    agent = MockAgent()
    tm = TaskManager(agent, bus)

    t1 = await tm.submit("Task 1")
    t2 = await tm.submit("Task 2")
    await tm.execute(t1.id)

    pending = tm.list_tasks(status=TaskStatus.PENDING)
    completed = tm.list_tasks(status=TaskStatus.COMPLETED)

    assert len(pending) == 1
    assert pending[0].id == t2.id
    assert len(completed) == 1
    assert completed[0].id == t1.id


async def test_running_count(bus: EventBus):
    import asyncio

    agent = MockAgent(delay=0.05)
    tm = TaskManager(agent, bus)

    task = await tm.submit("Slow task")

    # Start execution in background
    exec_task = asyncio.create_task(tm.execute(task.id))
    await asyncio.sleep(0.01)

    assert tm.running_count == 1

    await exec_task
    assert tm.running_count == 0


async def test_submit_emits_event(bus: EventBus):
    received: list[Event] = []

    async def handler(event: Event):
        received.append(event)

    bus.subscribe(handler, EventType.TASK_CREATED)

    agent = MockAgent()
    tm = TaskManager(agent, bus)
    task = await tm.submit("Test task", priority=5)

    assert len(received) == 1
    assert received[0].data["task_id"] == task.id
    assert received[0].data["priority"] == 5
