"""TaskGraph runner."""

import inspect
import threading
from concurrent.futures import Executor
from enum import Enum, auto
from typing import Any

from abacus_common_logic.concurrent.task_graph import (
    Task,
    TaskGraph,
    TaskGraphError,
    TaskOutput,
)
from abacus_common_logic.utils.logging import task_id_var


class TaskStatus(Enum):
    """Task execution states."""

    FAILED = auto()
    SKIPPED = auto()
    SUCCESS = auto()


class TaskGraphRunner:
    """Executes tasks in a DAG with parallel execution and dependency management.

    Manages the execution lifecycle of tasks in a directed acyclic graph,
    handling task scheduling, dependency resolution, result propagation,
    and error handling. Uses a thread pool for parallel execution.
    """

    def __init__(
        self,
        graph: TaskGraph,
        pools: dict[str, Executor],
        default_pool: str | None = None,
    ):
        """Initialize the task graph runner.

        Args:
            graph: TaskGraph object defining the task dependencies.
            pools: Map of available pools.
            default_pool: Default pool if no pool specified by the task.
                If not given, a default pool is chosen by the implementation.

        Raises:
            ValueError: If pools is empty.
        """
        if not pools:
            raise ValueError('At least one pool must be provided.')

        self.graph = graph
        self._pools = pools
        self._default_pool = (
            default_pool if default_pool else next(iter(self._pools.keys()))
        )

        # Properties
        self._first_error: Exception | None = None
        self._has_run = False
        self._in_degrees = graph.in_degrees.copy()
        self._results: dict[str, Any] = {}
        self._condition = threading.Condition()

        # State
        self._tasks_pending = set(graph.tasks.keys())
        self._tasks_running: set[str] = set()
        self._tasks_complete: set[str] = set()
        self._tasks_failed: set[str] = set()
        self._tasks_skipped: set[str] = set()

    def run(self, timeout: float | None = None) -> dict[str, Any]:
        """Execute all tasks in the graph according to their dependencies.

        Tasks are executed in topological order with parallel execution where
        dependencies allow. Execution stops on first error, skipping remaining tasks.

        Args:
            timeout: Optional timeout in seconds for the entire workflow.

        Raises:
            RuntimeError: If the runner is reused or if the workflow fails.
            TimeoutError: If the workflow exceeds the timeout.
        """
        if self._has_run:
            raise RuntimeError('Runner cannot be reused.')
        self._has_run = True

        try:
            self._execute_graph(timeout)
            if self._first_error:
                raise RuntimeError(
                    f'Workflow failed: {self._first_error}'
                ) from self._first_error
            return self._results
        finally:
            self._cleanup()

    def _execute_graph(self, timeout: float | None) -> None:
        """Execute the graph traversal logic.

        Args:
            timeout: Optional timeout in seconds.
        """
        for tid, deg in self._in_degrees.items():
            if deg == 0:
                self._submit(tid)

        self._wait_for_completion(timeout)

    def _cleanup(self) -> None:
        """Private cleanup to release memory and resources."""
        with self._condition:
            self._in_degrees.clear()

    def _on_task_complete(self, task_id: str, status: TaskStatus, result: Any) -> None:
        """Handle task completion and trigger dependent tasks.

        Updates task tracking sets, stores results, decrements dependency counts
        for downstream tasks, and submits newly-ready tasks for execution.

        Args:
            task_id: ID of the completed task.
            status: Final status of the task (SUCCESS, FAILED, or SKIPPED).
            result: Return value from the task handler (if successful).
        """
        with self._condition:
            try:
                self._tasks_running.discard(task_id)

                if status == TaskStatus.SUCCESS:
                    self._results[task_id] = result
                    self._tasks_complete.add(task_id)

                    if not self._first_error:
                        for neighbor in self.graph.adjacency_list[task_id]:
                            self._in_degrees[neighbor] -= 1
                            if self._in_degrees[neighbor] == 0:
                                self._submit(neighbor)
                elif status == TaskStatus.SKIPPED:
                    self._tasks_skipped.add(task_id)
                elif status == TaskStatus.FAILED:
                    self._tasks_failed.add(task_id)
                else:
                    raise ValueError(f"Invalid status '{status}' for task '{task_id}'")
            finally:
                self._condition.notify_all()

    def _resolve_params(self, params: dict[str, Any]) -> dict[str, Any]:
        """Resolve TaskOutput placeholders to actual task results.

        Replaces TaskOutput objects in params with the actual output from
        completed tasks, optionally extracting nested values via path.

        Args:
            params: Parameter dictionary that may contain TaskOutput placeholders.

        Returns:
            dict with TaskOutput placeholders replaced by actual values.
        """
        final_params: dict[str, Any] = {}

        for k, v in params.items():
            if isinstance(v, TaskOutput):
                val = self._results[v.task_id]
                final_params[k] = self._resolve_path(val, v.path) if v.path else val
            else:
                final_params[k] = v

        return final_params

    def _resolve_path(self, obj: Any, path: str) -> Any:
        """Extract nested value from object using dot-separated path.

        Supports dictionary keys, list/tuple indices, and object attributes.

        Args:
            obj: Object to extract value from.
            path: Dot-separated path (e.g., 'data.items.0.name').

        Returns:
            Value at the specified path.

        Raises:
            KeyError: If dictionary key doesn't exist.
            IndexError: If list index is out of range.
            AttributeError: If object attribute doesn't exist.
        """
        current = obj
        for part in path.split('.'):
            if isinstance(current, dict):
                current = current[part]
            elif isinstance(current, (list, tuple)):
                current = current[int(part)]
            else:
                current = getattr(current, part)
        return current

    def _run_task(self, task: Task) -> None:
        """Execute a task with error handling and result storage.

        Resolves TaskOutput placeholders to actual values, executes the task
        handler, stores results, and triggers callbacks. On failure, either
        calls the task's error handler or aborts the workflow.

        Args:
            task: Task object to execute.
        """
        status = TaskStatus.FAILED
        result = None
        token = task_id_var.set(task.id)

        try:
            with self._condition:
                self._tasks_running.add(task.id)
                self._tasks_pending.discard(task.id)

            if self._first_error:
                status = TaskStatus.SKIPPED
                return

            params = self._resolve_params(task.params)
            if 'executor' not in params:
                sig = inspect.signature(task.handler)
                if 'executor' in sig.parameters:
                    params['executor'] = self._try_pool(task.pool)
            result = task.handler(**params)
            status = TaskStatus.SUCCESS
        except Exception as e:
            with self._condition:
                if self._first_error is None:
                    self._first_error = TaskGraphError(task.id, e)
        finally:
            task_id_var.reset(token)
            self._on_task_complete(task.id, status, result)

    def _submit(self, task_id: str) -> None:
        """Submit a task to the thread pool for execution.

        Args:
            task_id: ID of the task to submit.
        """
        if self._first_error:
            return
        task = self.graph.tasks[task_id]
        self._try_pool(task.pool).submit(self._run_task, task)

    def _try_pool(self, pool_key: str | None) -> Executor:
        """Get the executor pool by key, or the default pool if key is None.

        Args:
            pool_key: The key of the pool to retrieve.

        Returns:
            Executor: The requested thread pool.

        Raises:
            ValueError: If the pool key does not exist.
        """
        executor = self._pools.get(pool_key or self._default_pool)
        if not executor:
            raise ValueError(f"Pool with key '{pool_key}' does not exist.")
        return executor

    def _wait_for_completion(self, timeout: float | None) -> None:
        """Wait for all pending and running tasks to complete or timeout.

        Blocks until either all tasks finish or the timeout expires.
        Sets _first_error to TimeoutError if timeout occurs.

        Args:
            timeout: Maximum seconds to wait (None for indefinite wait).

        Raises:
            TimeoutError: If timeout expires before all tasks complete.
        """
        with self._condition:
            finished = self._condition.wait_for(
                lambda: (
                    (len(self._tasks_pending) == 0 and len(self._tasks_running) == 0)
                    or self._first_error
                ),
                timeout=timeout,
            )
            if not finished:
                self._first_error = TimeoutError(f'Workflow timed out after {timeout}s')
                raise self._first_error