"""SSH connection configuration and ControlMaster pooling.

Centralises all SSH flags so proxy.py and sync.py don't duplicate
connection options. Uses ControlMaster for connection multiplexing —
first connection ~1.5s, subsequent ~50-160ms.
"""

from __future__ import annotations

import contextlib
import logging
import signal
import subprocess
import sys
from pathlib import Path

from specter_cli.config import specter_dir

log = logging.getLogger(__name__)


def _control_dir() -> Path:
    """ControlMaster socket directory — one socket per host."""
    return specter_dir() / "ssh"


def _ensure_control_dir() -> Path:
    """Create the ControlMaster socket directory if it doesn't exist."""
    d = _control_dir()
    d.mkdir(parents=True, exist_ok=True)
    # SSH requires 700 permissions on the control socket directory.
    d.chmod(0o700)
    return d


def _control_path(host: str) -> str:
    """Return the ControlMaster socket path for a given host."""
    control_dir = _ensure_control_dir()
    return str(control_dir / "%r@%h:%p")


def ssh_base_args(
    host: str,
    user: str,
    *,
    port: int = 22,
    key_path: str | None = None,
    tty: bool = False,
) -> list[str]:
    """Build base SSH argument list with ControlMaster options.

    Parameters
    ----------
    host:
        Remote hostname or IP.
    user:
        SSH username.
    port:
        SSH port.
    key_path:
        Path to SSH private key. If None, SSH uses its defaults.
    tty:
        If True, allocate a PTY (-t). Required for interactive sessions,
        but merges stdout/stderr.

    Returns
    -------
    List of SSH arguments (without the leading ``ssh`` command).
    """
    args: list[str] = []

    # ControlMaster for connection reuse.
    control = _control_path(host)
    args.extend(
        [
            "-o",
            "ControlMaster=auto",
            "-o",
            f"ControlPath={control}",
            "-o",
            "ControlPersist=600",
        ]
    )

    # Standard options.
    args.extend(
        [
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "UserKnownHostsFile=/dev/null",
            "-o",
            "LogLevel=ERROR",
        ]
    )

    if port != 22:
        args.extend(["-p", str(port)])

    if key_path:
        args.extend(["-i", key_path])

    if tty:
        args.append("-t")

    args.append(f"{user}@{host}")
    return args


def rsync_ssh_args(
    *,
    port: int = 22,
    key_path: str | None = None,
) -> str:
    """Build the SSH command string for rsync's ``-e`` flag.

    Returns a string like ``"ssh -o ControlMaster=auto -p 2222"``.
    This is passed directly to ``rsync -e '...'``.
    """
    parts = ["ssh"]

    # ControlMaster for rsync too.
    control_dir = _ensure_control_dir()
    control = str(control_dir / "%r@%h:%p")
    parts.extend(
        [
            "-o",
            "ControlMaster=auto",
            "-o",
            f"ControlPath={control}",
            "-o",
            "ControlPersist=600",
        ]
    )

    parts.extend(
        [
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "UserKnownHostsFile=/dev/null",
            "-o",
            "LogLevel=ERROR",
        ]
    )

    if port != 22:
        parts.extend(["-p", str(port)])

    if key_path:
        parts.extend(["-i", key_path])

    return " ".join(parts)


def check_connection(
    host: str,
    user: str,
    *,
    port: int = 22,
    key_path: str | None = None,
    timeout: int = 10,
) -> bool:
    """Test SSH connectivity. Returns True if connection succeeds."""
    args = ["ssh", *ssh_base_args(host, user, port=port, key_path=key_path), "echo ok"]
    try:
        result = subprocess.run(
            args,
            capture_output=True,
            text=True,
            timeout=timeout,
        )
        return result.returncode == 0
    except (subprocess.TimeoutExpired, OSError):
        return False


def stop_connection(
    host: str,
    user: str,
    *,
    port: int = 22,
) -> None:
    """Tear down the ControlMaster connection for a host."""
    control = _control_path(host)
    args = [
        "ssh",
        "-o",
        f"ControlPath={control}",
        "-O",
        "exit",
        f"{user}@{host}",
    ]
    if port != 22:
        args.extend(["-p", str(port)])

    with contextlib.suppress(subprocess.TimeoutExpired, OSError):
        subprocess.run(args, capture_output=True, text=True, timeout=5)


def exec_ssh_passthrough(cmd: list[str]) -> int:
    """Run a command with stdin/stdout/stderr passthrough and signal forwarding.

    Used by proxy.py and forward.py to execute SSH commands with full I/O
    passthrough.  Handles SIGINT/SIGTERM forwarding to the child process
    and restores original signal handlers on completion.

    Parameters
    ----------
    cmd:
        The full command to execute (e.g. ``["ssh", "-o", "...", "root@host", "nvidia-smi"]``).

    Returns
    -------
    The process exit code (1 if returncode is None, i.e. process was killed).
    """
    proc = subprocess.Popen(
        cmd,
        stdin=sys.stdin,
        stdout=sys.stdout,
        stderr=sys.stderr,
        start_new_session=True,
    )

    def _forward_signal(signum: int, _frame: object) -> None:
        """Forward signal to the subprocess."""
        with contextlib.suppress(ProcessLookupError):
            proc.send_signal(signum)

    prev_int = signal.signal(signal.SIGINT, _forward_signal)
    prev_term = signal.signal(signal.SIGTERM, _forward_signal)

    try:
        proc.wait()
    finally:
        signal.signal(signal.SIGINT, prev_int)
        signal.signal(signal.SIGTERM, prev_term)

    return proc.returncode if proc.returncode is not None else 1
