from __future__ import annotations

import shlex
import subprocess
from collections.abc import Sequence
from dataclasses import dataclass


@dataclass
class CollectResult:
    """Raw collection output for one host.

    Attributes:
        host: Host name used for collection.
        stdout: Command standard output.
        stderr: Command standard error.
        error: Normalized error message when collection failed; otherwise None.
    """

    host: str
    stdout: str
    stderr: str
    error: str | None = None


@dataclass(frozen=True)
class SSHOptions:
    """Optional SSH transport tuning options."""

    use_mux: bool = False
    control_persist_seconds: int = 30
    control_path: str | None = "~/.ssh/hytop-%C"
    server_alive_interval: int = 5
    server_alive_count_max: int = 1


DEFAULT_SSH_OPTIONS = SSHOptions(
    use_mux=True,
    control_persist_seconds=30,
    control_path="~/.ssh/hytop-%C",
    server_alive_interval=5,
    server_alive_count_max=1,
)


def _build_ssh_option_args(ssh_timeout: float, ssh_options: SSHOptions | None) -> list[str]:
    effective = ssh_options or DEFAULT_SSH_OPTIONS
    connect_timeout = max(1, round(ssh_timeout))
    options = [
        "-o",
        "BatchMode=yes",
        "-o",
        f"ConnectTimeout={connect_timeout}",
        "-o",
        f"ServerAliveInterval={max(1, effective.server_alive_interval)}",
        "-o",
        f"ServerAliveCountMax={max(1, effective.server_alive_count_max)}",
    ]
    if effective.use_mux:
        options.extend(["-o", "ControlMaster=auto"])
        options.extend(["-o", f"ControlPersist={max(1, effective.control_persist_seconds)}s"])
        if effective.control_path:
            options.extend(["-o", f"ControlPath={effective.control_path}"])
    return options


def _build_remote_python_shell_command(python_code: str) -> str:
    """Build one safely-quoted remote shell command for python -c."""

    return f"python3 -c {shlex.quote(python_code)}"


def collect_from_host(
    host: str,
    ssh_timeout: float,
    cmd_timeout: float,
    hy_smi_args: Sequence[str],
    ssh_options: SSHOptions | None = None,
) -> CollectResult:
    """Run hy-smi locally or via SSH and return raw output.

    Args:
        host: Hostname or localhost alias.
        ssh_timeout: SSH connect timeout in seconds.
        cmd_timeout: Command timeout in seconds.

    Returns:
        Raw command output with normalized error information.
    """
    local_names = {"localhost", "127.0.0.1", "::1"}
    if host in local_names:
        cmd = ["hy-smi", *hy_smi_args]
    else:
        cmd = [
            "ssh",
            *_build_ssh_option_args(ssh_timeout=ssh_timeout, ssh_options=ssh_options),
            host,
            "hy-smi",
            *hy_smi_args,
        ]

    try:
        proc = subprocess.run(
            cmd,
            check=False,
            capture_output=True,
            text=True,
            timeout=cmd_timeout,
        )
    except subprocess.TimeoutExpired:
        return CollectResult(
            host=host,
            stdout="",
            stderr="",
            error=f"timeout after {cmd_timeout:.1f}s",
        )
    except OSError as exc:
        return CollectResult(host=host, stdout="", stderr="", error=str(exc))

    if proc.returncode != 0:
        stderr = proc.stderr.strip() or "unknown error"
        return CollectResult(
            host=host,
            stdout=proc.stdout,
            stderr=proc.stderr,
            error=f"exit {proc.returncode}: {stderr}",
        )

    return CollectResult(host=host, stdout=proc.stdout, stderr=proc.stderr, error=None)


def collect_python_from_host(
    host: str,
    ssh_timeout: float,
    cmd_timeout: float,
    python_code: str,
    ssh_options: SSHOptions | None = None,
) -> CollectResult:
    """Run Python code locally or via SSH and return raw output."""

    local_names = {"localhost", "127.0.0.1", "::1"}
    if host in local_names:
        cmd = ["python3", "-c", python_code]
    else:
        cmd = [
            "ssh",
            *_build_ssh_option_args(ssh_timeout=ssh_timeout, ssh_options=ssh_options),
            host,
            _build_remote_python_shell_command(python_code),
        ]

    try:
        proc = subprocess.run(
            cmd,
            check=False,
            capture_output=True,
            text=True,
            timeout=cmd_timeout,
        )
    except subprocess.TimeoutExpired:
        return CollectResult(
            host=host,
            stdout="",
            stderr="",
            error=f"timeout after {cmd_timeout:.1f}s",
        )
    except OSError as exc:
        return CollectResult(host=host, stdout="", stderr="", error=str(exc))

    if proc.returncode != 0:
        stderr = proc.stderr.strip() or "unknown error"
        return CollectResult(
            host=host,
            stdout=proc.stdout,
            stderr=proc.stderr,
            error=f"exit {proc.returncode}: {stderr}",
        )

    return CollectResult(host=host, stdout=proc.stdout, stderr=proc.stderr, error=None)


def build_remote_python_command(
    host: str,
    ssh_timeout: float,
    python_code: str,
    ssh_options: SSHOptions | None = None,
) -> list[str]:
    """Build command for remote Python execution."""

    return [
        "ssh",
        *_build_ssh_option_args(ssh_timeout=ssh_timeout, ssh_options=ssh_options),
        host,
        _build_remote_python_shell_command(python_code),
    ]
