"""QueueManager: decides which task to dispatch next.

Responsibilities:
- Priority-ordered dispatch (URGENT > HIGH > NORMAL > LOW)
- Per-project lock (prevents concurrent agents on the same git repo)
- Scheduling: cron (via croniter), interval (in-memory last_run), after (dependency)
"""

from __future__ import annotations

from datetime import UTC, datetime, timedelta

from croniter import croniter

from dobby.models import ScheduleType, Task, TaskStatus
from dobby.store.interface import DobbyStore


class QueueManager:
    """Wraps a DobbyStore and controls task dispatch order."""

    def __init__(self, store: DobbyStore) -> None:
        self.store = store
        self._locked_projects: set[str] = set()
        self._running_counts: dict[str, int] = {}  # project -> running count

    # ------------------------------------------------------------------
    # Project locking
    # ------------------------------------------------------------------

    def lock_project(self, name: str) -> None:
        """Mark a project as having an active agent (no new tasks dispatched)."""
        self._locked_projects.add(name)
        self._running_counts[name] = self._running_counts.get(name, 0) + 1

    def unlock_project(self, name: str) -> None:
        """Release the project lock when the agent finishes."""
        count = self._running_counts.get(name, 0) - 1
        if count <= 0:
            self._locked_projects.discard(name)
            self._running_counts.pop(name, None)
        else:
            self._running_counts[name] = count

    def is_locked(self, name: str) -> bool:
        return name in self._locked_projects

    def running_count(self, name: str) -> int:
        return self._running_counts.get(name, 0)

    async def _can_dispatch_to_project(self, project_name: str) -> bool:
        """Check if a project can accept another task dispatch.

        Uses max_concurrent from project policies if set, otherwise
        falls back to the default single-task-per-project lock.
        """
        proj = await self.store.get_project(project_name)
        max_conc = None
        if proj is not None and proj.policies.max_concurrent is not None:
            max_conc = proj.policies.max_concurrent

        current = self._running_counts.get(project_name, 0)
        if max_conc is not None:
            return current < max_conc
        # Default: single concurrency (locked means no more)
        return project_name not in self._locked_projects

    # ------------------------------------------------------------------
    # Next task selection
    # ------------------------------------------------------------------

    async def next_task(self) -> Task | None:
        """Return the highest-priority PENDING task whose project is not locked.

        Returns None if there are no eligible tasks.
        """
        # get_pending_tasks returns tasks sorted by priority DESC, created_at ASC
        tasks = await self.store.get_pending_tasks(limit=100)
        for task in tasks:
            if await self._can_dispatch_to_project(task.project):
                return task
        return None

    async def next_dispatchable_task(self) -> Task | None:
        """Return the highest-priority unlocked task that can run now.

        Includes both pending one-off tasks and active recurring tasks that are due.
        """
        pending = await self.store.get_pending_tasks(limit=100)
        due_scheduled = await self.get_due_scheduled_tasks()

        all_candidates = pending + due_scheduled
        all_candidates.sort(
            key=lambda task: (-task.priority.weight, task.created_at)
        )
        for task in all_candidates:
            # Skip tasks targeted at specific spokes — those are claimed via the spoke API
            target = getattr(task, 'target', 'local')
            if target not in ('local', 'any', '', None):
                continue
            if await self._can_dispatch_to_project(task.project):
                return task
        return None

    # ------------------------------------------------------------------
    # Scheduling
    # ------------------------------------------------------------------

    async def get_due_scheduled_tasks(self) -> list[Task]:
        """Return scheduled tasks (CRON / INTERVAL) that are due to run now.

        Uses pre-computed next_run_at if available; falls back to computing
        from croniter / interval math and persists the computed next_run_at.
        """
        now = datetime.now(UTC)
        scheduled = await self.store.get_scheduled_tasks()
        due: list[Task] = []

        for task in scheduled:
            sched = task.schedule

            if sched.type == ScheduleType.CRON:
                if sched.expr is None:
                    continue
                next_dt = task.next_run_at
                if next_dt is None:
                    next_dt = self.compute_next_run(task)
                    if next_dt is not None:
                        await self.store.update_task(
                            task.id, next_run_at=next_dt.isoformat(),
                        )
                if next_dt is not None:
                    if next_dt.tzinfo is None:
                        next_dt = next_dt.replace(tzinfo=UTC)
                    if next_dt <= now:
                        due.append(task)

            elif sched.type == ScheduleType.INTERVAL:
                next_dt = task.next_run_at
                if next_dt is None:
                    next_dt = self.compute_next_run(task)
                    if next_dt is not None:
                        await self.store.update_task(
                            task.id, next_run_at=next_dt.isoformat(),
                        )
                if next_dt is None:
                    # Never run and no next_run_at — due immediately
                    due.append(task)
                else:
                    if next_dt.tzinfo is None:
                        next_dt = next_dt.replace(tzinfo=UTC)
                    if next_dt <= now:
                        due.append(task)

        return due

    @staticmethod
    def compute_next_run(task: Task) -> datetime | None:
        """Compute the next execution time for a recurring task."""
        sched = task.schedule
        if sched.type == ScheduleType.CRON:
            if sched.expr is None:
                return None
            last = task.last_run_at or task.created_at
            if last.tzinfo is None:
                last = last.replace(tzinfo=UTC)
            cron = croniter(sched.expr, last)
            next_dt = cron.get_next(datetime)
            if next_dt.tzinfo is None:
                next_dt = next_dt.replace(tzinfo=UTC)
            return next_dt
        elif sched.type == ScheduleType.INTERVAL:
            interval_secs = sched.interval_seconds()
            last = task.last_run_at
            if last is None:
                return None  # Due immediately — no next_run_at needed
            if last.tzinfo is None:
                last = last.replace(tzinfo=UTC)
            return last + timedelta(seconds=interval_secs)
        return None

    # ------------------------------------------------------------------
    # Dependency resolution
    # ------------------------------------------------------------------

    async def check_dependencies(self) -> None:
        """Resolve AFTER-type scheduled tasks.

        If the task they depend on is COMPLETED, transition them to PENDING
        so they become eligible for dispatch.
        """
        # Get all tasks that have an AFTER schedule and are not yet PENDING/COMPLETED
        all_tasks = await self.store.list_tasks()
        for task in all_tasks:
            if task.schedule.type != ScheduleType.AFTER:
                continue
            dep_id = task.schedule.depends_on
            if dep_id is None:
                continue
            # Skip tasks that are already pending or completed
            if task.status in (TaskStatus.PENDING, TaskStatus.COMPLETED, TaskStatus.CANCELLED):
                continue
            try:
                dep = await self.store.get_task(dep_id)
            except KeyError:
                continue
            if dep.status == TaskStatus.COMPLETED:
                await self.store.update_task(task.id, status=TaskStatus.PENDING.value)
