"""Task abstraction for pipeline DAG nodes.

This module provides the Task class which wraps callables as nodes in a
directed acyclic graph (DAG) for pipeline execution.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional


@dataclass
class TaskResult:
    """Result of executing a pipeline task.

    Attributes:
        data: The primary output data from the task. For pipeline steps, this
            typically contains a StepResult object. Access via task_result.data
            to get the StepResult.
        metadata: Dictionary of metadata about the task execution (e.g., timing,
            parameters used, etc.).
        files_generated: List of file paths generated by this task.
    """

    data: Any
    metadata: Dict[str, Any] = field(default_factory=dict)
    files_generated: List[str | Path] = field(default_factory=list)


class Task:
    """A node in the pipeline DAG.

    Each task wraps a callable function and declares its dependencies on other
    tasks. The task name is used to reference it in the DAG and must be unique
    within a pipeline.

    Args:
        func: The callable to execute. Should accept keyword arguments matching
            the outputs of its dependencies.
        name: Unique identifier for this task within the pipeline.
        depends_on: List of task names that must execute before this task.
        output_prefix: Prefix to use for output files generated by this task.
        description: Human-readable description of what this task does.

    Example:
        >>> def load_data(data_path: str) -> pd.DataFrame:
        ...     return pd.read_csv(data_path)
        >>>
        >>> def clean_data(data: pd.DataFrame) -> pd.DataFrame:
        ...     return data.dropna()
        >>>
        >>> load_task = Task(
        ...     func=load_data,
        ...     name="load_data",
        ...     depends_on=[],
        ...     description="Load raw data from CSV"
        ... )
        >>>
        >>> clean_task = Task(
        ...     func=clean_data,
        ...     name="clean_data",
        ...     depends_on=["load_data"],
        ...     description="Remove missing values"
        ... )
    """

    def __init__(
        self,
        func: Callable,
        name: str,
        depends_on: Optional[List[str]] = None,
        output_prefix: str = "",
        description: str = "",
    ):
        """Initialize a Task.

        Args:
            func: The callable to execute.
            name: Unique identifier for this task.
            depends_on: List of task names this task depends on.
            output_prefix: Prefix for output files.
            description: Description of the task.
        """
        self.func = func
        self.name = name
        self.depends_on = depends_on or []
        self.output_prefix = output_prefix
        self.description = description

    def __repr__(self) -> str:
        """Return string representation of the task."""
        deps = ", ".join(self.depends_on) if self.depends_on else "none"
        return f"Task(name='{self.name}', depends_on=[{deps}])"

    def execute(
        self,
        config: Any,
        run_dir: Path,
        logger: Any,
        dependency_results: Optional[Dict[str, TaskResult]] = None,
    ) -> TaskResult:
        """Execute the task function.

        Args:
            config: Configuration object for the pipeline.
            run_dir: Directory for this pipeline run's outputs.
            logger: Logger instance for recording execution.
            dependency_results: Dictionary mapping dependency task names to their
                TaskResults. Used to pass outputs from dependencies to this task.

        Returns:
            TaskResult containing the task's output data, metadata, and generated files.

        Raises:
            Exception: Any exception raised by the task function is propagated.
        """
        import time

        logger.info(f"Executing task: {self.name}")

        # Prepare kwargs from dependency results
        kwargs = {"config": config, "run_dir": run_dir, "logger": logger}

        if dependency_results:
            for dep_name in self.depends_on:
                if dep_name in dependency_results:
                    # Pass the entire TaskResult from the dependency
                    # Tasks can access .data, .metadata, .files_generated as needed
                    kwargs[dep_name] = dependency_results[dep_name]

        # Execute the task
        start_time = time.time()
        try:
            result = self.func(**kwargs)
            elapsed = time.time() - start_time

            # If the function returns a TaskResult, use it directly
            if isinstance(result, TaskResult):
                result.metadata["elapsed_time"] = elapsed
                result.metadata["task_name"] = self.name
                logger.info(
                    f"Task {self.name} completed in {elapsed:.2f}s, "
                    f"generated {len(result.files_generated)} files"
                )
                return result

            # Otherwise, wrap the result in a TaskResult
            task_result = TaskResult(
                data=result,
                metadata={
                    "elapsed_time": elapsed,
                    "task_name": self.name,
                },
                files_generated=[],
            )
            logger.info(f"Task {self.name} completed in {elapsed:.2f}s")
            return task_result

        except Exception as e:
            elapsed = time.time() - start_time
            logger.error(f"Task {self.name} failed after {elapsed:.2f}s: {e}")
            raise
