"""Task lifecycle management with parallel execution support."""

from __future__ import annotations

import asyncio
import time

from agent_scout.agent import Agent
from agent_scout.db.models import Task, TaskStatus
from agent_scout.events import Event, EventBus, EventType


class TaskManager:
    """Manages task lifecycle, queue, and parallel execution.

    Features:
    - Priority-based task queue
    - Configurable concurrency limit for parallel execution
    - Task dependencies (one task can block another)
    - Full lifecycle tracking
    """

    def __init__(
        self,
        agent: Agent,
        event_bus: EventBus,
        max_concurrent: int = 3,
    ) -> None:
        self._agent = agent
        self._bus = event_bus
        self._max_concurrent = max_concurrent
        self._tasks: dict[str, Task] = {}
        self._dependencies: dict[str, set[str]] = {}  # task_id -> set of blocking task_ids
        self._running: dict[str, asyncio.Task[Task]] = {}
        self._semaphore = asyncio.Semaphore(max_concurrent)

    async def submit(
        self,
        description: str,
        priority: int = 0,
        budget_limit: float | None = None,
        depends_on: list[str] | None = None,
    ) -> Task:
        """Submit a new task to the manager.

        The task is queued and will be executed when resources are available
        and all dependencies are satisfied.
        """
        task = Task(
            description=description,
            priority=priority,
            budget_limit=budget_limit,
        )
        self._tasks[task.id] = task

        if depends_on:
            self._dependencies[task.id] = set(depends_on)

        await self._bus.emit(
            Event(
                event_type=EventType.TASK_CREATED,
                data={
                    "task_id": task.id,
                    "description": description,
                    "priority": priority,
                },
                task_id=task.id,
            )
        )

        return task

    async def execute(self, task_id: str) -> Task:
        """Execute a single task, respecting concurrency limits."""
        task = self._tasks.get(task_id)
        if not task:
            raise ValueError(f"Task {task_id} not found")

        # Wait for dependencies
        await self._wait_for_dependencies(task_id)

        # Respect concurrency limit
        async with self._semaphore:
            asyncio_task = asyncio.create_task(self._agent.run_task(task))
            self._running[task_id] = asyncio_task
            try:
                result = await asyncio_task
                self._tasks[task_id] = result
                return result
            finally:
                self._running.pop(task_id, None)

    async def execute_all(self) -> list[Task]:
        """Execute all pending tasks, respecting priorities and dependencies.

        Tasks run in parallel up to the concurrency limit.
        """
        pending = self._get_pending_sorted()
        if not pending:
            return []

        coros = [self.execute(t.id) for t in pending]
        results = await asyncio.gather(*coros, return_exceptions=True)

        completed = []
        for result in results:
            if isinstance(result, Task):
                completed.append(result)
            elif isinstance(result, Exception):
                # Task failed — already handled by agent
                pass
        return completed

    async def cancel(self, task_id: str) -> None:
        """Cancel a pending or running task."""
        task = self._tasks.get(task_id)
        if not task:
            return

        # If running, stop the agent loop
        if task_id in self._running:
            await self._agent.stop_task(task_id)
            self._running[task_id].cancel()
            self._running.pop(task_id, None)

        task.status = TaskStatus.CANCELLED
        task.updated_at = time.time()

        await self._bus.emit(
            Event(
                event_type=EventType.TASK_STATE_CHANGE,
                data={"status": "cancelled"},
                task_id=task_id,
            )
        )

    def get_task(self, task_id: str) -> Task | None:
        return self._tasks.get(task_id)

    def list_tasks(self, status: TaskStatus | None = None) -> list[Task]:
        tasks = list(self._tasks.values())
        if status:
            tasks = [t for t in tasks if t.status == status]
        return sorted(tasks, key=lambda t: (-t.priority, t.created_at))

    @property
    def running_count(self) -> int:
        return len(self._running)

    @property
    def pending_count(self) -> int:
        return sum(1 for t in self._tasks.values() if t.status == TaskStatus.PENDING)

    async def _wait_for_dependencies(self, task_id: str) -> None:
        """Wait until all dependencies of a task are completed."""
        deps = self._dependencies.get(task_id, set())
        while deps:
            remaining = set()
            for dep_id in deps:
                dep_task = self._tasks.get(dep_id)
                if dep_task and dep_task.status not in (
                    TaskStatus.COMPLETED,
                    TaskStatus.FAILED,
                    TaskStatus.CANCELLED,
                ):
                    remaining.add(dep_id)
            if not remaining:
                break
            await asyncio.sleep(0.01)
            deps = remaining

    def _get_pending_sorted(self) -> list[Task]:
        """Get pending tasks sorted by priority (highest first), then creation time."""
        return sorted(
            (t for t in self._tasks.values() if t.status == TaskStatus.PENDING),
            key=lambda t: (-t.priority, t.created_at),
        )
