"""DAG-based task graph."""

from dataclasses import dataclass, field
from typing import (
    Any,
    Callable,
    Generic,
    ParamSpec,
    TypeVar,
)

P = ParamSpec('P')
R = TypeVar('R')


@dataclass
class TaskOutput:
    """Placeholder for task output that will be resolved at execution time.

    Used in task parameters to reference the output of another task.

    Attributes:
        task_id: ID of the task whose output this placeholder represents.

    Example:
        Task(id="process", handler=process_data, params={"input": TaskOutput("fetch")})
    """

    task_id: str
    path: str | None = None


@dataclass(frozen=True)
class Task(Generic[P, R]):
    """A single executable task in a DAG workflow.

    Attributes:
        id: Unique identifier for the task.
        handler: Callable function to execute.
        params: Parameters to pass to the handler.
        depends_on: Set of task IDs that must complete before this task.
    """

    id: str
    handler: Callable[P, R]
    params: dict[str, Any] = field(default_factory=dict, compare=False)
    depends_on: set[str] = field(default_factory=set, compare=False)
    pool: str | None = None


class TaskGraph:
    """Manages the structural integrity of the DAG."""

    def __init__(self, tasks: list[Task]):
        """Initialize the task graph with validation.

        Args:
            tasks: List of Task objects to build the graph from.

        Raises:
            ValueError: If the graph contains cycles or invalid data flow.
        """
        self.tasks = self._build_task_map(tasks)
        self.adjacency_list = self._build_adjacency_list(tasks)
        self.in_degrees = {t.id: len(t.depends_on) for t in tasks}
        self._validate_cycles()
        self._validate_data_flow()

    def _build_adjacency_list(self, tasks: list[Task]) -> dict[str, list[str]]:
        """Build adjacency list representation of task dependencies.

        Args:
            tasks: List of Task objects to build graph from.

        Returns:
            dict mapping task IDs to list of dependent task IDs.

        Raises:
            ValueError: If a task depends on an unknown task.
        """
        adj: dict[str, list[str]] = {task.id: [] for task in tasks}

        for task in tasks:
            for dep in task.depends_on:
                if dep not in adj:
                    raise ValueError(
                        f"Task '{task.id}' depends on unknown task '{dep}'"
                    )
                adj[dep].append(task.id)

        return adj

    def _build_task_map(self, tasks: list[Task]) -> dict[str, Task]:
        """Build a dictionary of tasks keyed by ID.

        Args:
            tasks: List of Task objects.

        Returns:
            dict: Mapping of task ID to Task object.

        Raises:
            ValueError: If duplicate task IDs are found.
        """
        dupes: set[str] = set()
        task_map: dict[str, Task] = {}

        for task in tasks:
            if task.id in task_map:
                dupes.add(task.id)
            else:
                task_map[task.id] = task

        if dupes:
            raise ValueError(f'Duplicate task IDs found: {", ".join(dupes)}')

        return task_map

    def _validate_cycles(self) -> None:
        """Validate that the task graph contains no cycles using Kahn's Algorithm.

        Raises:
            ValueError: If a cycle is detected in the task dependencies.
        """
        # Kahn's Algorithm
        in_degrees = self.in_degrees.copy()

        # Get tasks with zero pending dependencies
        queue = [tid for tid, deg in in_degrees.items() if deg == 0]
        visited = 0

        # For each task in the queue
        while queue:
            u = queue.pop(0)
            visited += 1
            # Update each dependent's dependency count
            for v in self.adjacency_list.get(u, []):
                in_degrees[v] -= 1
                # Add to the queue if no pending dependencies
                if in_degrees[v] == 0:
                    queue.append(v)

        # Check if all tasks can be visited
        if visited == len(self.tasks):
            return

        # Cycle(s) detected
        unvisited = [t_id for t_id, deg in in_degrees.items() if deg > 0]
        raise ValueError(f'Graph cycle(s) detected for tasks: {unvisited}')

    def _validate_data_flow(self) -> None:
        """Validate that TaskOutput references only point to ancestor tasks.

        Ensures data dependencies are valid - a task can only consume output
        from tasks that execute before it in the DAG.

        Raises:
            ValueError: If a task references output from a non-ancestor task.
        """
        # Ancestry check for TaskOutput usage
        for task in self.tasks.values():
            for val in task.params.values():
                if not isinstance(val, TaskOutput):
                    continue
                target = val.task_id
                if target not in self.tasks:
                    raise ValueError(
                        f"Task '{task.id}' requests output from "
                        f"non-existent task '{target}'."
                    )
                if not self.is_ancestor(val.task_id, task.id):
                    raise ValueError(
                        f"Task '{task.id}' requests output from "
                        f"non-ancestor task '{target}'. "
                        f'This would cause a race condition.'
                    )

    def is_ancestor(self, target: str, current: str) -> bool:
        """Check if target_id is an ancestor of current_id.

        Args:
            target: The potential ancestor task ID.
            current: The descendant task ID.

        Returns:
            bool: True if target is an ancestor of current, False otherwise.
        """
        stack = list(self.tasks[current].depends_on)
        visited = set()

        while stack:
            dep = stack.pop()
            if dep == target:
                return True
            if dep not in visited:
                visited.add(dep)
                stack.extend(self.tasks[dep].depends_on)

        return False


class TaskGraphError(Exception):
    """Exception raised when a task in the graph fails during execution."""

    def __init__(self, task_id: str, original_error: Exception):
        """Initialize the task graph error.

        Args:
            task_id: ID of the task that failed.
            original_error: The original exception that caused the failure.
        """
        self.task_id = task_id
        super().__init__(f"Task '{task_id}' failed: {original_error}")