"""SSH config management, port helpers, and SSH/SCP wrappers."""

from __future__ import annotations

import re
import socket
import subprocess
import sys
from pathlib import Path

import vastly
from vastly import dim

SSH_CONFIG_DIR = Path.home() / ".ssh" / "vast.d"

_GENERATED_HEADER = "# Generated by vastly -- do not edit"


def _atomic_write(path: Path, content: str) -> None:
    """Write content to path atomically via temp-file-then-rename."""
    tmp = path.parent / (path.name + ".vastly-tmp")
    tmp.write_text(content, encoding="utf-8")
    try:
        tmp.replace(path)
    except OSError:
        # Windows: replace() fails if target is open; fall back to delete+rename
        path.unlink(missing_ok=True)
        tmp.rename(path)


# Quick probes (echo ok, test -f, find): fail fast on dead connections
SSH_OPTS = [
    "-o",
    "ConnectTimeout=10",
    "-o",
    "ServerAliveInterval=5",
    "-o",
    "ServerAliveCountMax=2",
]

# Long-running commands (setup script): allow more time between keepalives
SSH_SETUP_OPTS = [
    "-o",
    "ConnectTimeout=10",
    "-o",
    "ServerAliveInterval=15",
    "-o",
    "ServerAliveCountMax=4",
]


def is_port_available(port: int) -> bool:
    """Check if a local TCP port is available to bind."""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        try:
            sock.bind(("127.0.0.1", port))
            return True
        except OSError:
            return False


def find_available_port(start: int, exclude: set[int] | None = None) -> int:
    """Find the first available port starting from *start*, skipping *exclude*."""
    exclude = exclude or set()
    port = start
    while port <= 65535 and (not is_port_available(port) or port in exclude):
        port += 1
    if port > 65535:
        raise RuntimeError(f"No available port found starting from {start}")
    return port


def ensure_ssh_include() -> None:
    """Ensure ~/.ssh/config contains 'Include vast.d/*'."""
    ssh_dir = Path.home() / ".ssh"
    ssh_config = ssh_dir / "config"
    include_line = "Include vast.d/*"

    if ssh_config.exists():
        content = ssh_config.read_text(encoding="utf-8")
        if not re.search(r"Include\s+vast\.d[/\\]\*", content):
            _atomic_write(ssh_config, f"{include_line}\n{content}")
            print(dim(f"Added '{include_line}' to {ssh_config}"))
    else:
        ssh_dir.mkdir(parents=True, exist_ok=True)
        _atomic_write(ssh_config, include_line + "\n")
        print(dim(f"Created {ssh_config} with '{include_line}'"))


def write_ssh_config(
    name: str,
    *,
    host: str,
    port: int,
    user: str,
    key_path: str | None,
    local_forwards: list[tuple[int, int]],
) -> None:
    """Write a single SSH config file to ~/.ssh/vast.d/<name>."""
    SSH_CONFIG_DIR.mkdir(parents=True, exist_ok=True)

    lines = [
        _GENERATED_HEADER,
        "# Vast.ai instances are ephemeral with recycled IPs, so host key",
        "# verification is disabled. Agent forwarding enables git clone via SSH.",
        f"Host {name}",
        f"    HostName {host}",
        f"    Port {port}",
        f"    User {user}",
        "    ForwardAgent yes",
        "    StrictHostKeyChecking no",
        f"    UserKnownHostsFile {'NUL' if sys.platform == 'win32' else '/dev/null'}",
        "    LogLevel ERROR",
    ]

    if key_path:
        lines.append(f"    IdentityFile {key_path}")

    for local_port, remote_port in local_forwards:
        lines.append(f"    LocalForward {local_port} localhost:{remote_port}")

    config_file = SSH_CONFIG_DIR / name
    _atomic_write(config_file, "\n".join(lines) + "\n")


def clear_ssh_configs() -> None:
    """Remove vastly-generated files from ~/.ssh/vast.d/."""
    if SSH_CONFIG_DIR.exists():
        for f in SSH_CONFIG_DIR.iterdir():
            if not f.is_file():
                continue
            first_line = f.read_text(encoding="utf-8").split("\n", 1)[0]
            if first_line.startswith("# Generated by vastly"):
                f.unlink()


def prune_ssh_configs(keep: set[str]) -> None:
    """Remove vastly-generated configs from ~/.ssh/vast.d/ that aren't in *keep*."""
    if SSH_CONFIG_DIR.exists():
        for f in SSH_CONFIG_DIR.iterdir():
            if not f.is_file():
                continue
            if f.name in keep:
                continue
            first_line = f.read_text(encoding="utf-8").split("\n", 1)[0]
            if first_line.startswith("# Generated by vastly"):
                f.unlink()


def run_ssh(
    host: str, command: str, *, setup: bool = False, stream: bool = False
) -> subprocess.CompletedProcess[str]:
    """Run an SSH command on *host*.

    Args:
        setup: Use longer timeout options (for setup scripts).
        stream: Print output in real time instead of capturing.
    """
    opts = list(SSH_SETUP_OPTS if setup else SSH_OPTS)
    if vastly.VERBOSE:
        opts.extend(["-o", "LogLevel=DEBUG"])
    cmd = ["ssh", *opts, host, command]

    vastly.verbose(f"ssh {host}: {command[:80]}")

    # Hard timeout: SSH options cover TCP connect and keepalives, but not
    # handshake stalls.  Probe commands get 30s; setup scripts get 10min.
    deadline = 600 if setup else 30

    try:
        if stream:
            return subprocess.run(cmd, text=True, timeout=deadline)
        return subprocess.run(cmd, capture_output=True, text=True, timeout=deadline)
    except subprocess.TimeoutExpired as e:
        return subprocess.CompletedProcess(
            cmd,
            returncode=1,
            stdout=e.stdout or "",
            stderr=e.stderr or "timeout",
        )


def run_scp(
    src: str,
    dest: str,
    *,
    setup: bool = False,
    recursive: bool = False,
    stream: bool = False,
) -> subprocess.CompletedProcess[str]:
    """SCP a file from *src* to *dest*.

    Args:
        setup: Use longer timeout options (for setup scripts).
        recursive: Copy directories recursively.
        stream: Show SCP progress instead of capturing output.
    """
    opts = list(SSH_SETUP_OPTS if setup else SSH_OPTS)
    if vastly.VERBOSE:
        opts.extend(["-o", "LogLevel=DEBUG"])
    flags = ["-r"] if recursive else []
    deadline = 600 if setup else 60

    vastly.verbose(f"scp {src} -> {dest}")

    try:
        if stream:
            return subprocess.run(
                ["scp", *flags, *opts, src, dest], text=True, timeout=deadline
            )
        return subprocess.run(
            ["scp", *flags, *opts, src, dest],
            capture_output=True,
            text=True,
            timeout=deadline,
        )
    except subprocess.TimeoutExpired as e:
        return subprocess.CompletedProcess(
            ["scp", src, dest],
            returncode=1,
            stdout=e.stdout or "",
            stderr=e.stderr or "timeout",
        )
