"""Stub file for sky.utils.command_runner.

This file is dynamically generated by stubgen and added with the
overloaded type hints for SSHCommandRunner.run(), as we need to
determine the return type based on the value of require_outputs.
"""
import enum
import typing
from typing import Any, Callable, Iterable, List, Optional, Tuple, Union

from typing_extensions import Literal

from sky import sky_logging as sky_logging
from sky.skylet import log_lib as log_lib
from sky.utils import subprocess_utils as subprocess_utils

GIT_EXCLUDE: str
RSYNC_DISPLAY_OPTION: str
RSYNC_FILTER_GITIGNORE: str
RSYNC_FILTER_SKYIGNORE: str
RSYNC_EXCLUDE_OPTION: str
ALIAS_SUDO_TO_EMPTY_FOR_ROOT_CMD: str
DEFAULT_SSH_CONTROL_NAME: str


def ssh_options_list(
    ssh_private_key: Optional[str],
    ssh_control_name: Optional[str],
    *,
    ssh_proxy_command: Optional[str] = ...,
    ssh_proxy_jump: Optional[str] = ...,
    docker_ssh_proxy_command: Optional[str] = ...,
    connect_timeout: Optional[int] = ...,
    port: int = ...,
    disable_control_master: Optional[bool] = ...,
    escape_percent_expand: bool = ...,
    ssh_log_file: Optional[str] = ...,
    disable_identities_only: bool = ...,
) -> List[str]:
    ...


class SshMode(enum.Enum):
    NON_INTERACTIVE = ...
    INTERACTIVE = ...
    LOGIN = ...


class CommandRunner:
    node_id: str

    def __init__(
        self,
        node: Tuple[Any, ...],
        **kwargs,
    ) -> None:
        ...

    @typing.overload
    def run(self,
            cmd: Union[str, List[str]],
            *,
            require_outputs: Literal[False] = ...,
            log_path: str = ...,
            process_stream: bool = ...,
            stream_logs: bool = ...,
            separate_stderr: bool = ...,
            connect_timeout: Optional[int] = ...,
            source_bashrc: bool = ...,
            skip_lines: int = ...,
            run_in_background: bool = ...,
            **kwargs) -> int:
        ...

    @typing.overload
    def run(self,
            cmd: Union[str, List[str]],
            *,
            require_outputs: Literal[True],
            log_path: str = ...,
            process_stream: bool = ...,
            stream_logs: bool = ...,
            separate_stderr: bool = ...,
            connect_timeout: Optional[int] = ...,
            source_bashrc: bool = ...,
            skip_lines: int = ...,
            run_in_background: bool = ...,
            **kwargs) -> Tuple[int, str, str]:
        ...

    @typing.overload
    def run(self,
            cmd: Union[str, List[str]],
            *,
            require_outputs: bool = ...,
            log_path: str = ...,
            process_stream: bool = ...,
            stream_logs: bool = ...,
            separate_stderr: bool = ...,
            connect_timeout: Optional[int] = ...,
            source_bashrc: bool = ...,
            skip_lines: int = ...,
            run_in_background: bool = ...,
            **kwargs) -> Union[Tuple[int, str, str], int]:
        ...

    def get_remote_home_dir(self) -> str:
        ...

    def rsync(
        self,
        source: str,
        target: str,
        *,
        up: bool,
        log_path: str = ...,
        stream_logs: bool = ...,
        max_retry: int = ...,
    ) -> None:
        ...

    @typing.overload
    def run_driver(self,
                   cmd: Union[str, List[str]],
                   *,
                   require_outputs: Literal[False] = ...,
                   **kwargs) -> int:
        ...

    @typing.overload
    def run_driver(self, cmd: Union[str, List[str]], *,
                   require_outputs: Literal[True],
                   **kwargs) -> Tuple[int, str, str]:
        ...

    @typing.overload
    def run_driver(self,
                   cmd: Union[str, List[str]],
                   *,
                   require_outputs: bool = ...,
                   **kwargs) -> Union[Tuple[int, str, str], int]:
        ...

    @typing.overload
    def run_setup(self,
                  cmd: Union[str, List[str]],
                  *,
                  require_outputs: Literal[False] = ...,
                  **kwargs) -> int:
        ...

    @typing.overload
    def run_setup(self, cmd: Union[str, List[str]], *,
                  require_outputs: Literal[True],
                  **kwargs) -> Tuple[int, str, str]:
        ...

    @typing.overload
    def run_setup(self,
                  cmd: Union[str, List[str]],
                  *,
                  require_outputs: bool = ...,
                  **kwargs) -> Union[Tuple[int, str, str], int]:
        ...

    def rsync_driver(self,
                     source: str,
                     target: str,
                     *,
                     up: bool,
                     log_path: str = ...,
                     stream_logs: bool = ...,
                     max_retry: int = ...) -> None:
        ...

    def rsync_setup(self,
                    source: str,
                    target: str,
                    *,
                    up: bool,
                    log_path: str = ...,
                    stream_logs: bool = ...,
                    max_retry: int = ...) -> None:
        ...

    def port_forward_command(
            self,
            port_forward: List[Tuple[int, int]],
            connect_timeout: int = 1,
            ssh_mode: SshMode = SshMode.INTERACTIVE) -> List[str]:
        ...

    @classmethod
    def make_runner_list(cls: typing.Type[CommandRunner],
                         node_list: Iterable[Tuple[Any, ...]],
                         **kwargs) -> List[CommandRunner]:
        ...

    def check_connection(self) -> bool:
        ...

    def close_cached_connection(self) -> None:
        ...


class SSHCommandRunner(CommandRunner):
    ip: str
    port: int
    ssh_user: str
    ssh_private_key: Optional[str]
    ssh_control_name: Optional[str]
    docker_user: str
    disable_control_master: Optional[bool]
    port_forward_execute_remote_command: Optional[bool]
    enable_interactive_auth: bool
    disable_identities_only: bool

    def __init__(
        self,
        node: Tuple[str, int],
        ssh_user: str,
        ssh_private_key: Optional[str],
        ssh_control_name: Optional[str] = ...,
        ssh_proxy_command: Optional[str] = ...,
        ssh_proxy_jump: Optional[str] = ...,
        docker_user: Optional[str] = ...,
        disable_control_master: Optional[bool] = ...,
        port_forward_execute_remote_command: Optional[bool] = ...,
        enable_interactive_auth: bool = ...,
        disable_identities_only: bool = ...,
    ) -> None:
        ...

    @typing.overload
    def run(self,
            cmd: Union[str, List[str]],
            *,
            require_outputs: Literal[False] = ...,
            port_forward: Optional[List[int]] = ...,
            log_path: str = ...,
            process_stream: bool = ...,
            stream_logs: bool = ...,
            ssh_mode: SshMode = ...,
            separate_stderr: bool = ...,
            connect_timeout: Optional[int] = ...,
            source_bashrc: bool = ...,
            skip_lines: int = ...,
            run_in_background: bool = ...,
            **kwargs) -> int:
        ...

    @typing.overload
    def run(self,
            cmd: Union[str, List[str]],
            *,
            require_outputs: Literal[True],
            port_forward: Optional[List[int]] = ...,
            log_path: str = ...,
            process_stream: bool = ...,
            stream_logs: bool = ...,
            ssh_mode: SshMode = ...,
            separate_stderr: bool = ...,
            connect_timeout: Optional[int] = ...,
            source_bashrc: bool = ...,
            skip_lines: int = ...,
            run_in_background: bool = ...,
            **kwargs) -> Tuple[int, str, str]:
        ...

    @typing.overload
    def run(self,
            cmd: Union[str, List[str]],
            *,
            require_outputs: bool = ...,
            port_forward: Optional[List[int]] = ...,
            log_path: str = ...,
            process_stream: bool = ...,
            stream_logs: bool = ...,
            ssh_mode: SshMode = ...,
            separate_stderr: bool = ...,
            connect_timeout: Optional[int] = ...,
            source_bashrc: bool = ...,
            skip_lines: int = ...,
            run_in_background: bool = ...,
            **kwargs) -> Union[Tuple[int, str, str], int]:
        ...

    def ssh_base_command(
        self,
        *,
        ssh_mode: SshMode,
        port_forward: Optional[List[Tuple[int, int]]],
        connect_timeout: Optional[int],
    ) -> List[str]:
        ...

    def rsync(
        self,
        source: str,
        target: str,
        *,
        up: bool,
        log_path: str = ...,
        stream_logs: bool = ...,
        max_retry: int = ...,
        get_remote_home_dir: Callable[[], str] = ...,
    ) -> None:
        ...

    def port_forward_command(
            self,
            port_forward: List[Tuple[int, int]],
            connect_timeout: int = 1,
            ssh_mode: SshMode = SshMode.INTERACTIVE) -> List[str]:
        ...


class KubernetesCommandRunner(CommandRunner):

    def __init__(
        self,
        node: Tuple[Tuple[str, Optional[str]], str],
        deployment: Optional[str] = ...,
        container: Optional[str] = ...,
        **kwargs,
    ) -> None:
        ...

    @typing.overload
    def run(self,
            cmd: Union[str, List[str]],
            *,
            port_forward: Optional[List[int]] = ...,
            require_outputs: Literal[False] = ...,
            log_path: str = ...,
            process_stream: bool = ...,
            stream_logs: bool = ...,
            ssh_mode: SshMode = ...,
            separate_stderr: bool = ...,
            connect_timeout: Optional[int] = ...,
            source_bashrc: bool = ...,
            skip_lines: int = ...,
            run_in_background: bool = ...,
            **kwargs) -> int:
        ...

    @typing.overload
    def run(self,
            cmd: Union[str, List[str]],
            *,
            port_forward: Optional[List[int]] = ...,
            require_outputs: Literal[True],
            log_path: str = ...,
            process_stream: bool = ...,
            stream_logs: bool = ...,
            ssh_mode: SshMode = ...,
            separate_stderr: bool = ...,
            connect_timeout: Optional[int] = ...,
            source_bashrc: bool = ...,
            skip_lines: int = ...,
            run_in_background: bool = ...,
            **kwargs) -> Tuple[int, str, str]:
        ...

    @typing.overload
    def run(self,
            cmd: Union[str, List[str]],
            *,
            port_forward: Optional[List[int]] = ...,
            require_outputs: bool = ...,
            log_path: str = ...,
            process_stream: bool = ...,
            stream_logs: bool = ...,
            ssh_mode: SshMode = ...,
            separate_stderr: bool = ...,
            connect_timeout: Optional[int] = ...,
            source_bashrc: bool = ...,
            skip_lines: int = ...,
            run_in_background: bool = ...,
            **kwargs) -> Union[Tuple[int, str, str], int]:
        ...

    def rsync(
        self,
        source: str,
        target: str,
        *,
        up: bool,
        log_path: str = ...,
        stream_logs: bool = ...,
        max_retry: int = ...,
    ) -> None:
        ...

    def port_forward_command(
            self,
            port_forward: List[Tuple[int, int]],
            connect_timeout: int = 1,
            ssh_mode: SshMode = SshMode.INTERACTIVE) -> List[str]:
        ...


class SlurmCommandRunner(SSHCommandRunner):
    """Runner for Slurm commands."""
    sky_dir: str
    skypilot_runtime_dir: str
    job_id: str
    slurm_node: str
    container_args: Optional[str]

    def __init__(
        self,
        node: Tuple[str, int],
        ssh_user: str,
        ssh_private_key: Optional[str],
        *,
        sky_dir: str,
        skypilot_runtime_dir: str,
        job_id: str,
        slurm_node: str,
        container_args: Optional[str] = ...,
        **kwargs,
    ) -> None:
        ...


class LocalProcessCommandRunner(CommandRunner):

    def __init__(self) -> None:
        ...

    @typing.overload
    def run(self,
            cmd: Union[str, List[str]],
            *,
            port_forward: Optional[List[int]] = ...,
            require_outputs: Literal[False] = ...,
            log_path: str = ...,
            process_stream: bool = ...,
            stream_logs: bool = ...,
            ssh_mode: SshMode = ...,
            separate_stderr: bool = ...,
            connect_timeout: Optional[int] = ...,
            source_bashrc: bool = ...,
            skip_lines: int = ...,
            run_in_background: bool = ...,
            **kwargs) -> int:
        ...

    @typing.overload
    def run(self,
            cmd: Union[str, List[str]],
            *,
            port_forward: Optional[List[int]] = ...,
            require_outputs: Literal[True],
            log_path: str = ...,
            process_stream: bool = ...,
            stream_logs: bool = ...,
            ssh_mode: SshMode = ...,
            separate_stderr: bool = ...,
            connect_timeout: Optional[int] = ...,
            source_bashrc: bool = ...,
            skip_lines: int = ...,
            run_in_background: bool = ...,
            **kwargs) -> Tuple[int, str, str]:
        ...

    @typing.overload
    def run(self,
            cmd: Union[str, List[str]],
            *,
            port_forward: Optional[List[int]] = ...,
            require_outputs: bool = ...,
            log_path: str = ...,
            process_stream: bool = ...,
            stream_logs: bool = ...,
            ssh_mode: SshMode = ...,
            separate_stderr: bool = ...,
            connect_timeout: Optional[int] = ...,
            source_bashrc: bool = ...,
            skip_lines: int = ...,
            run_in_background: bool = ...,
            **kwargs) -> Union[Tuple[int, str, str], int]:
        ...
