"""AI Infrastructure Service CLI - SSH."""

import json
import os
from enum import Enum
from typing import Annotated

import typer
from slices_clientlib_aiis.schema.job import Job, SSHInfo

from . import console, err_console
from .ai import ai_app
from .ai_common import get_ai_job
from .ai_context import AIContext
from .ai_helpers import (
    handle_some_exceptions,
    test_ssh_connect,
)


def _unnest_jump_proxies(first_jump_proxy: SSHInfo | None) -> list[SSHInfo]:
    """Return a list of all jump proxies, instead of the nested representation.

    This is easier to use in code.
    The list is also reversed, that is, the jump proxy all others depend on, comes first.
    """

    res = []
    # The loop below should always be finite.
    # But it is always best to have a safety.
    # Because infinite loops cause end user bugs that are hard to debug.
    safety = 5
    jump_proxy = first_jump_proxy
    while jump_proxy:
        res.insert(0, jump_proxy)
        jump_proxy = jump_proxy.jump_proxy
        safety -= 1
        if safety <= 0:
            err_console.print("[red]BUG: Too much loops handling jump proxy :x:")
            err_console.print("[red]Please report this bug!")
            raise typer.Exit(code=1)

    return res


def _make_jump_proxy_cli_args_and_config(unnested_jump_proxies: list[SSHInfo]) -> tuple[list[str], str]:
    ssh_jump_proxy_destinations = []
    for jump_proxy in unnested_jump_proxies:
        port_dest = f":{jump_proxy.port}" if jump_proxy.port != DEFAULT_SSH_PORT else ""
        ssh_jump_proxy_destinations.append(f"{jump_proxy.username}@{jump_proxy.host}{port_dest}")

    return ["-J", ",".join(ssh_jump_proxy_destinations)], f"JumpProxy {','.join(ssh_jump_proxy_destinations)}"


def _get_jump_ssh_arg(jump: SSHInfo) -> list[str]:
    jump_proxies = _unnest_jump_proxies(jump)
    jump_proxy_cli_args, _ = _make_jump_proxy_cli_args_and_config(jump_proxies)
    return jump_proxy_cli_args


def _print_ssh_login(ssh_login: SSHInfo):
    params = []

    if ssh_login.port != DEFAULT_SSH_PORT:
        params.extend(["-p", f"{ssh_login.port}"])
    if ssh_login.jump_proxy:
        params.extend(_get_jump_ssh_arg(ssh_login.jump_proxy))

    return f"ssh{(' ' + ' '.join(params)) if params else ''} {ssh_login.username}@{ssh_login.host}"


class ProxyOption(Enum):
    """Proxy Options."""

    off = "off"
    on = "on"
    auto = "auto"


class ShowOption(Enum):
    """Show options."""

    nothing = "nothing"
    json = "json"
    command = "command"
    ansible = "ansible"
    ssh_config = "ssh_config"
    server_host = "server_host"
    proxy_host = "proxy_host"
    # server_pubkey_openssh = "server_pubkey_openssh"
    # proxy_pubkey_openssh = "proxy_pubkey_openssh"


DEFAULT_SSH_PORT = 22


def _job_to_ansible(job: Job, ssh_login: SSHInfo, proxy: bool):
    """Convert a resource to an Ansible inventory entry."""
    if proxy and ssh_login.jump_proxy:
        jump_proxies = _unnest_jump_proxies(ssh_login.jump_proxy)
        jump_proxy_cli_args, _ = _make_jump_proxy_cli_args_and_config(jump_proxies)
        jump_proxy_cli_args_str = " ".join(jump_proxy_cli_args)
        ansible_ssh_common_args = f" ansible_ssh_common_args='{jump_proxy_cli_args_str}'"
    else:
        ansible_ssh_common_args = ""

    return (
        f"gpulab-{job.short_id} "
        f"ansible_ssh_port={ssh_login.port} "
        f"ansible_ssh_host={ssh_login.host} "
        f"ansible_ssh_user={ssh_login.username} {ansible_ssh_common_args}"
    )


def _ssh_common(
    ctx: AIContext,
    job_id: str,
    proxy: ProxyOption,
    *,
    sftp_mode: bool = False,  # whether this is for sftp
) -> tuple[Job, SSHInfo, list[str], str, str]:
    job = get_ai_job(ctx, job_id)

    if not job.ssh_info:
        err_console.print("The resource has no SSH logins")
        raise typer.Exit(code=1)

    ssh_login = job.ssh_info[0]

    destination = f"{ssh_login.username}@{ssh_login.host}"

    if proxy == ProxyOption.auto:
        can_connect = test_ssh_connect(ssh_login.host, ssh_login.port)
        proxy = ProxyOption.on if can_connect else ProxyOption.off

        # method from legacy GPULab CLI
        # # Try to resolve worker, if it doesn't assume we need the proxy
        # # (reason: IDLab Gent hostnames have public IPv6, but not IPv4.)
        # try:
        #     host_ip = socket.gethostbyname(job.state.resources.slave_host)
        #     proxy = ProxyOption.off
        #
        #     # Also hardcode that uantwerpen needs proxy.
        #     # (reason: uantwerpen nodes are firewalled)
        #     if job.state.resources.slave_host.endswith(".idlab.uantwerpen.be"):
        #         proxy = ProxyOption.on
        # # except socket.gaierror as e:
        # except OSError as e:
        #     proxy = ProxyOption.on

    ssh_command_options: list[str] = []
    ssh_config_jump_proxy = ""

    if ssh_login.port != DEFAULT_SSH_PORT:
        port_option = "-p" if not sftp_mode else "-P"
        ssh_command_options.append(f"{port_option} {ssh_login.port}")

    jump_proxies = []
    if proxy == ProxyOption.on and ssh_login.jump_proxy:
        jump_proxies = _unnest_jump_proxies(ssh_login.jump_proxy)
        jump_proxy_cli_args, ssh_config_jump_proxy = _make_jump_proxy_cli_args_and_config(jump_proxies)
        ssh_command_options.extend(jump_proxy_cli_args)

    return job, ssh_login, ssh_command_options, destination, ssh_config_jump_proxy


def make_ssh_command(
    ctx: AIContext,
    job_id: str,
    proxy: ProxyOption = ProxyOption.auto,
) -> list[str]:
    """Make the command to connect to the ssh."""
    job, ssh_login, ssh_command_options, destination, ssh_config_jump_proxy = _ssh_common(ctx, job_id, proxy)

    return ["ssh", *ssh_command_options, destination]


@ai_app.command(
    "ssh",
    context_settings={"allow_extra_args": True, "ignore_unknown_options": True},
    rich_help_panel="Job interaction commands",
)
@handle_some_exceptions
def ssh_to_job(
    ctx: AIContext,
    job_id: Annotated[str, typer.Argument(help="Job ID.", show_default=False)],
    exec: Annotated[
        bool,
        typer.Option(
            "--exec/--no-exec",
            help="Exec SSH to connect to the job. (Default: --exec)",
            show_default=False,
        ),
    ] = True,
    show: Annotated[
        ShowOption,
        typer.Option(
            help="What to show: Nothing, SSH Command, Ansible inventory line, "
            "lines for ~/.ssh/config, or the server public key",
            show_default=False,
        ),
    ] = "command",  # pyright: ignore [reportArgumentType]  weird but this works
    proxy: Annotated[ProxyOption, typer.Option(help="Use an SSH jump proxy.", case_sensitive=False)] = "auto",  # pyright: ignore [reportArgumentType]  weird but this works
):
    """Connect to a job using SSH.

    Alternatively, show info, such as the login command.

    Extra arguments are passed directly to ssh.
    """
    # if show in (ShowOption.proxy_host, ShowOption.proxy_pubkey_openssh) and proxy == ProxyOption.auto:
    if show in (ShowOption.proxy_host,) and proxy == ProxyOption.auto:
        # if proxy info is requested, turn on proxy
        proxy = ProxyOption.on

    job, ssh_login, ssh_command_options, destination, ssh_config_jump_proxy = _ssh_common(ctx, job_id, proxy)

    command = ["ssh", *ssh_command_options, destination, *ctx.args]

    command_line = " ".join(command)

    do_exec_ssh = exec and show in [ShowOption.command, ShowOption.nothing]

    if show == ShowOption.command:
        if do_exec_ssh:
            console.print(f"[bold blue]{command_line}")
        else:
            typer.echo(command_line)  # output without any markup
    elif show == ShowOption.ansible:
        # output without any markup
        typer.echo(_job_to_ansible(job, ssh_login, proxy == ProxyOption.on))
    elif show == ShowOption.json:
        obj = {
            "host": ssh_login.host,
            "port": ssh_login.port,
            "username": ssh_login.username,
        }
        if ssh_login.jump_proxy:
            obj["jump_proxy"] = {
                "host": ssh_login.jump_proxy.host,
                "port": ssh_login.jump_proxy.port,
                "username": ssh_login.jump_proxy.username,
            }
        typer.echo(json.dumps(obj, indent=3))

    elif show == ShowOption.ssh_config:
        # output without any markup
        typer.echo(f"""Host={ssh_login.host}
   HostName {ssh_login.host}
   Port {ssh_login.port}
   User {ssh_login.username}
   PreferredAuthentications publickey
   IdentitiesOnly yes
   {ssh_config_jump_proxy}""")
    elif show == ShowOption.server_host:
        typer.echo(ssh_login.host)
    elif show == ShowOption.proxy_host:
        if ssh_login.jump_proxy:
            typer.echo(ssh_login.jump_proxy.host)
        else:
            err_console.print("No jump proxy")
    # elif show == ShowOption.server_pubkey_openssh:
    #     if ssh_login.host_keys:
    #         for host_key in ssh_login.host_keys:
    #             typer.echo(f"{host_key.strip()}")
    #     else:
    #         err_console.print("No hostkeys in SSH info")
    # elif show == ShowOption.proxy_pubkey_openssh:
    #     if ssh_login.jump_proxy:
    #         if ssh_login.jump_proxy.host_keys:
    #             for host_key in ssh_login.jump_proxy.host_keys:
    #                 typer.echo(f"{host_key.strip()}")
    #         else:
    #             err_console.print("No hostkeys in proxy SSH info")
    #     else:
    #         err_console.print("No jump proxy")

    if do_exec_ssh:
        os.execvp(command[0], command)  # noqa: S606


@ai_app.command(
    "sftp",
    context_settings={"allow_extra_args": True, "ignore_unknown_options": True},
    rich_help_panel="Job interaction commands",
)
@handle_some_exceptions
def sftp_to_resource(
    ctx: AIContext,
    job_id: Annotated[str, typer.Argument(help="ID of the job.", show_default=False)],
    exec: Annotated[
        bool,
        typer.Option(
            "--exec/--no-exec",
            help="Exec SFTP to connect to the resource. (Default: --exec)",
            show_default=False,
        ),
    ] = True,
    show_command: Annotated[
        bool,
        typer.Option(
            "--show-command/--no-show-command",
            help="Show command. (Default: --show-command)",
            show_default=False,
        ),
    ] = True,
    proxy: Annotated[ProxyOption, typer.Option(help="Use an SSH jump proxy.", case_sensitive=False)] = "auto",  # pyright: ignore [reportArgumentType]  weird but this works
):
    """Transfer files to/from a job using SFTP.

    The sftp tool is part of openssh and provides an interactive sftp shell.
    Extra arguments are passed directly to sftp.
    """
    job, ssh_login, ssh_command_options, destination, ssh_config_jump_proxy = _ssh_common(
        ctx, job_id, proxy, sftp_mode=True
    )

    command = ["sftp", *ssh_command_options, destination, *ctx.args]

    command_line = " ".join(command)

    if show_command:
        if exec:
            console.print(f"[bold blue]{command_line}")
        else:
            typer.echo(command_line)  # output without any markup

    if exec:
        os.execvp(command[0], command)  # noqa: S606
