"""Task execution protocol contract.

This module defines the Protocol that wt-task implements and wt-compiler
generates code against. This enables type-safe task execution without
direct dependencies between packages.
"""

from __future__ import annotations

from collections.abc import Callable, Sequence
from typing import TYPE_CHECKING, Any, ParamSpec, Protocol, TypeVar

if TYPE_CHECKING:
    from typing_extensions import Self

P = ParamSpec("P")
R = TypeVar("R", covariant=True)


class TaskProtocol(Protocol[P, R]):
    """Protocol defining the task execution interface.

    This protocol specifies the methods that must be implemented by task
    wrappers (e.g., wt-task's SyncTask). The wt-compiler generates code that
    calls these methods, ensuring type-safe execution without direct imports.

    Type parameters:
        P: Parameter specification for the wrapped function
        R: Return type of the wrapped function

    Required methods:
        partial: Partial function application
        call: Direct execution with arguments
        map: Map over a sequence of values for a single parameter
        mapvalues: Map over key-value pairs
        validate: Validate task configuration
        skipif: Conditional execution skipping
        set_executor: Set custom executor

    Examples:
        Type-checking a task implementation:

        >>> from typing import Callable, Any
        >>> def task_wrapper(func: Callable[P, R]) -> TaskProtocol[P, R]:
        ...     # Implementation would return SyncTask instance
        ...     ...

        Using in generated code (conceptual):

        >>> # Generated by wt-compiler:
        >>> # result = task(my_func).partial(x=1).map("y", [1, 2, 3])
    """

    def partial(self, **kwargs: Any) -> Self:
        """Apply partial function application.

        Binds keyword arguments to the task, returning a new task with those
        parameters fixed. This is used in generated DAG code to set parameters
        from spec.yaml.

        Args:
            **kwargs: Keyword arguments to bind

        Returns:
            New task with arguments partially applied

        Examples:
            >>> # In generated code:
            >>> # task(func).partial(x=1, y=2).call()
        """
        ...

    def call(self, *args: P.args, **kwargs: P.kwargs) -> R:
        """Execute the task with given arguments.

        Directly calls the wrapped function with provided arguments.

        Args:
            *args: Positional arguments
            **kwargs: Keyword arguments

        Returns:
            Result of function execution

        Examples:
            >>> # In generated code:
            >>> # result = task(func).call(x=1, y=2)
        """
        ...

    def map(self, argname: str, argvalues: Sequence[Any], **kwargs: Any) -> Sequence[R]:
        """Map over a sequence of values for a single parameter.

        Executes the task multiple times, once for each value in argvalues,
        binding that value to the parameter named argname. Other parameters
        can be provided via kwargs.

        Args:
            argname: Parameter name to map over
            argvalues: Sequence of values for that parameter
            **kwargs: Fixed keyword arguments for all executions

        Returns:
            Sequence of results, one per input value

        Examples:
            >>> # In generated code:
            >>> # results = task(func).map("x", [1, 2, 3], y=10)
            >>> # Equivalent to: [func(x=1, y=10), func(x=2, y=10), func(x=3, y=10)]
        """
        ...

    def mapvalues(
        self, argname: str, argvalues: Sequence[tuple[Any, Any]], **kwargs: Any
    ) -> Sequence[tuple[Any, R]]:
        """Map over key-value pairs, preserving keys.

        Similar to map, but input is a sequence of (key, value) tuples. Returns
        (key, result) tuples, preserving the keys. Useful for maintaining
        correspondence between inputs and outputs.

        Args:
            argname: Parameter name to map over
            argvalues: Sequence of (key, value) tuples
            **kwargs: Fixed keyword arguments for all executions

        Returns:
            Sequence of (key, result) tuples

        Examples:
            >>> # In generated code:
            >>> # results = task(func).mapvalues("x", [("a", 1), ("b", 2)], y=10)
            >>> # Returns: [("a", result1), ("b", result2)]
        """
        ...

    def validate(self) -> Self:
        """Validate task configuration.

        Performs validation checks on the task (e.g., Pydantic validation).
        Returns self for method chaining.

        Returns:
            Self for method chaining

        Examples:
            >>> # In generated code:
            >>> # task(func).partial(x=1).validate().call()
        """
        ...

    def skipif(self, condition: Callable[..., bool]) -> Self:
        """Conditionally skip task execution.

        Adds a condition that determines whether the task should be skipped.
        If condition returns True, task execution is skipped and a sentinel
        value is returned.

        Args:
            condition: Callable that returns True to skip, False to execute

        Returns:
            Self for method chaining

        Examples:
            >>> # In generated code:
            >>> # task(func).skipif(lambda: os.path.exists("skip.txt")).call()
        """
        ...

    def set_executor(self, executor: Any) -> Self:
        """Set custom executor for task execution.

        Allows specifying an alternative executor (e.g., thread pool, process
        pool, async executor) for this task.

        Args:
            executor: Executor instance

        Returns:
            Self for method chaining

        Examples:
            >>> # In generated code:
            >>> # from concurrent.futures import ThreadPoolExecutor
            >>> # executor = ThreadPoolExecutor(max_workers=4)
            >>> # task(func).set_executor(executor).map("x", [1, 2, 3])
        """
        ...
