"""Dependency resolution and ready queue management."""

from __future__ import annotations

import asyncpg
from redis.asyncio import Redis

from loom.exceptions import DependencyCycleError
from loom.graph import cache, store
from loom.graph.task import TaskStatus


async def check_and_unblock(
    pool: asyncpg.Pool,
    redis: Redis,
    completed_task_id: str,
    project_id: str,
) -> list[str]:
    """After a task completes, check dependents and unblock any that are now ready.

    Also handles epic auto-completion: if a completed task's parent is an epic
    and all its children are done, the epic is marked done too.

    Returns list of newly unblocked task IDs.
    """
    unblocked: list[str] = []
    dependent_ids = await store.get_dependents(pool, completed_task_id)

    if not dependent_ids:
        # Still check epic completion
        completed_task = await store.get_task(pool, completed_task_id)
        if completed_task.parent_id:
            await _check_epic_completion(pool, redis, completed_task.parent_id, project_id)
        return unblocked

    # Batch load all dependent tasks
    dep_tasks: dict[str, object] = {}
    for dep_id in dependent_ids:
        dep_tasks[dep_id] = await store.get_task(pool, dep_id)

    # Collect all dependency IDs we need to check
    all_req_ids: set[str] = set()
    for task in dep_tasks.values():
        if task.status == TaskStatus.BLOCKED:
            all_req_ids.update(task.depends_on)

    # Batch load their statuses
    req_statuses: dict[str, str] = {}
    if all_req_ids:
        async with pool.acquire() as conn:
            rows = await conn.fetch(
                "SELECT id, status FROM tasks WHERE id = ANY($1)",
                list(all_req_ids),
            )
        req_statuses = {r["id"]: r["status"] for r in rows}

    for dep_id, task in dep_tasks.items():
        if task.status != TaskStatus.BLOCKED:
            continue
        all_done = all(req_statuses.get(rid) == "done" for rid in task.depends_on)
        if all_done:
            updated = await store.set_task_status(pool, dep_id, TaskStatus.PENDING)
            await cache.sync_task(redis, updated)
            await cache.add_to_ready_queue(redis, updated)
            unblocked.append(dep_id)

    # Epic auto-completion: check if completed task's parent is an epic
    completed_task = await store.get_task(pool, completed_task_id)
    if completed_task.parent_id:
        await _check_epic_completion(pool, redis, completed_task.parent_id, project_id)

    return unblocked


async def _check_epic_completion(
    pool: asyncpg.Pool, redis: Redis, epic_id: str, project_id: str
) -> None:
    """If all children of an epic are done, mark the epic done.

    Walks up the parent chain so nested epics auto-close recursively.
    """
    epic = await store.get_task(pool, epic_id)
    if epic.status != TaskStatus.EPIC:
        return

    async with pool.acquire() as conn:
        children = await conn.fetch(
            "SELECT status FROM tasks WHERE parent_id = $1", epic_id
        )

    if not children:
        return

    if all(r["status"] == "done" for r in children):
        updated = await store.set_task_status(pool, epic_id, TaskStatus.DONE)
        await cache.sync_task(redis, updated)
        await cache.remove_from_ready_queue(redis, project_id, epic_id)
        # Recurse: if this epic has a parent, check if that parent is now complete too
        if updated.parent_id:
            await _check_epic_completion(pool, redis, updated.parent_id, project_id)


async def sweep_epic_completion(
    pool: asyncpg.Pool, redis: Redis, project_id: str,
) -> list[str]:
    """Scan all epics in a project and auto-close any whose children are all done.

    Returns list of epic IDs that were closed. Useful as a cleanup step
    when tasks were completed without parent_id links (e.g., manual loom_done).
    """
    closed: list[str] = []
    epics = await store.get_tasks_by_status(pool, project_id, "epic", limit=500)
    for epic in epics:
        async with pool.acquire() as conn:
            children = await conn.fetch(
                "SELECT status FROM tasks WHERE parent_id = $1", epic.id
            )
        if children and all(r["status"] == "done" for r in children):
            updated = await store.set_task_status(pool, epic.id, TaskStatus.DONE)
            await cache.sync_task(redis, updated)
            await cache.remove_from_ready_queue(redis, project_id, epic.id)
            closed.append(epic.id)
    return closed


async def detect_cycle(pool: asyncpg.Pool, task_id: str, new_deps: list[str]) -> None:
    """Check if adding new_deps to task_id would create a cycle. Raises DependencyCycleError."""
    # Iterative DFS from each new dep, following existing edges, looking for task_id
    for dep_id in new_deps:
        visited: set[str] = set()
        stack = [dep_id]
        while stack:
            current = stack.pop()
            if current == task_id:
                raise DependencyCycleError([task_id, dep_id, "...", task_id])
            if current in visited:
                continue
            visited.add(current)
            # Follow existing dependency edges
            async with pool.acquire() as conn:
                edges = await conn.fetch(
                    "SELECT depends_on FROM task_deps WHERE task_id = $1", current
                )
            for edge in edges:
                stack.append(edge["depends_on"])


async def compute_initial_status(
    pool: asyncpg.Pool, depends_on: list[str]
) -> TaskStatus:
    """Determine if a new task should start as pending or blocked."""
    if not depends_on:
        return TaskStatus.PENDING

    for dep_id in depends_on:
        dep = await store.get_task(pool, dep_id)
        if dep.status != TaskStatus.DONE:
            return TaskStatus.BLOCKED

    return TaskStatus.PENDING
