"""
Type stubs for core module
This file provides better IDE autocomplete and type checking
"""
from typing import Callable, TypeVar, Optional, Any, Protocol, overload, runtime_checkable, Type
from .models import ExecutionRun, TaskRun, StepRun, ExecutionStatus

T = TypeVar("T")

@runtime_checkable
class ExecutionFunction(Protocol[T]):
    """Protocol for execution functions with deploy method."""

    def __call__(self, *args: Any, **kwargs: Any) -> T: ...

    def deploy(
        self,
        instance_id: str,
        secret_key: str,
        server: str,
        port: int,
        host: str = '0.0.0.0',
        encryption_key: Optional[str] = None,
        **config: Any
    ) -> ExecutionFunction[T]:
        """Deploy this execution to the server for remote triggering.

        Args:
            instance_id (str): Instance ID for this robot deployment.
            secret_key (str): Secret key for platform authentication.
            server (str): Platform server URL for status updates.
            port (int): Port to run server on.
            host (str, optional): Host to bind server to. Defaults to '0.0.0.0'.
            encryption_key (str, optional): Encryption key for decrypting credentials.
            **config (Any): Additional configuration available to all tasks/steps.
        """
        ...

# Overloads for execution decorator
@overload
def execution(
    func: Callable[..., T],
) -> ExecutionFunction[T]: ...

@overload
def execution(
    *,
    name: Optional[str] = None,
    version: Optional[str] = None,
    description: Optional[str] = None,
    exception_handlers: Optional[dict[Type[Exception], ExecutionStatus]] = None,
) -> Callable[[Callable[..., T]], ExecutionFunction[T]]: ...

def execution(
    func: Optional[Callable[..., T]] = None,
    *,
    name: Optional[str] = None,
    version: Optional[str] = None,
    description: Optional[str] = None,
    exception_handlers: Optional[dict[Type[Exception], ExecutionStatus]] = None,
) -> Callable[[Callable[..., T]], ExecutionFunction[T]] | ExecutionFunction[T]:
    """Decorator to mark a function as an execution.

    Args:
        func (Callable, optional): The function to decorate.
        name (str, optional): Custom name for the execution.
        version (str, optional): Version of the execution.
        description (str, optional): Description of the execution.
        exception_handlers (dict, optional): Map of exception types to ExecutionStatus.
            When an exception occurs anywhere (step, task, or execution),
            the execution will be marked with the corresponding status.

    Usage:
        @execution(
            exception_handlers={
                KeyboardInterrupt: ExecutionStatus.INTERRUPTED,
                OutOfHoursError: ExecutionStatus.OUT_OF_OPERATING_HOURS
            }
        )
        def my_execution():
            pass
    """
    ...

@overload
def task(
    func: Callable[..., T],
) -> Callable[..., T]: ...

@overload
def task(
    *,
    name: Optional[str] = None,
    retries: int = 0,
    description: Optional[str] = None,
    code: Optional[str] = None,
    observation: Optional[str] = None,
) -> Callable[[Callable[..., T]], Callable[..., T]]: ...

def task(
    func: Optional[Callable[..., T]] = None,
    *,
    name: Optional[str] = None,
    retries: int = 0,
    description: Optional[str] = None,
    code: Optional[str] = None,
    observation: Optional[str] = None,
) -> Callable[[Callable[..., T]], Callable[..., T]] | Callable[..., T]:
    """Decorator to mark a function as a task.

    Args:
        func (Callable, optional): The function to decorate.
        name (str, optional): Custom name for the task.
        retries (int, optional): Number of retries on failure. Defaults to 0.
        description (str, optional): Description of the task.
        code (str, optional): Task code for platform identification.
        observation (str, optional): Initial observation message.
    """
    ...

@overload
def step(
    *,
    code: str,
    name: Optional[str] = None,
    description: Optional[str] = None,
    observation: Optional[str] = None,
) -> Callable[[Callable[..., T]], Callable[..., T]]: ...

def step(
    func: Optional[Callable[..., T]] = None,
    *,
    code: str,
    name: Optional[str] = None,
    description: Optional[str] = None,
    observation: Optional[str] = None,
) -> Callable[[Callable[..., T]], Callable[..., T]] | Callable[..., T]:
    """Decorator to mark a function as a step.

    Args:
        func (Callable, optional): The function to decorate.
        code (str): Step code for platform identification. Required.
        name (str, optional): Custom name for the step.
        description (str, optional): Description of the step.
        observation (str, optional): Initial observation message.
    """
    ...

def get_current_execution() -> Optional[ExecutionRun]: ...
def get_current_task() -> Optional[TaskRun]: ...
def set_total_tasks(count: int) -> None: ...
