"""SSH client wrapper for the dms-ui.

Provides a single class, :class:`DMSSSH`, that executes a shell command on a
remote DMS host and returns ``(stdout, stderr, exit_code)``. We deliberately
do **not** keep a long-lived connection between requests — every API call opens
a fresh connection, runs the command, and closes it. This keeps the layer
simple and avoids stuck/half-dead sessions.

Two backends are supported, tried in order:

1. ``paramiko`` if it is importable.
2. ``subprocess`` + ``ssh -o BatchMode=yes`` with ``SSH_ASKPASS`` pointed at
   ``/home/hermes/.ssh-askpass.sh`` — this is the same trick used by
   ``mcp_server.py`` for the sharekit project.

The active backend is recorded in :data:`BACKEND` for diagnostic logging.
"""
from __future__ import annotations

import logging
import os
import re
import shlex
import subprocess
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Tuple

log = logging.getLogger("dms-ui.ssh")

ASKPASS_PATH = Path("/home/hermes/.ssh-askpass.sh")
# SSH client config: point ssh at /home/hermes/.ssh/config so that
# aliases (e.g. "smail-server" → smail.icu) resolve correctly when dms-ui
# is launched by the root user (e.g. via systemd or /etc/nginx reverse proxy).
SSH_CONFIG_PATH = Path(os.environ.get("DMS_UI_SSH_CONFIG", "/home/hermes/.ssh/config"))
ANSI_RE = re.compile(rb"\x1b\[[0-9;?]*[A-Za-z]")
BACKEND: str = "unknown"
_PARAMIKO_IMPORT_ERROR: Optional[Exception] = None

try:
    import paramiko  # type: ignore

    BACKEND = "paramiko"
except Exception as _exc:  # pragma: no cover - import-time probe
    paramiko = None  # type: ignore
    _PARAMIKO_IMPORT_ERROR = _exc
    BACKEND = "subprocess"


@dataclass
class CommandResult:
    """The result of a remote command."""

    stdout: str
    stderr: str
    exit_code: int
    backend: str = BACKEND

    @property
    def ok(self) -> bool:
        return self.exit_code == 0

    def to_dict(self) -> dict:
        return {
            "ok": self.ok,
            "stdout": self.stdout,
            "stderr": self.stderr,
            "exit_code": self.exit_code,
            "backend": self.backend,
        }


def _strip_ansi(data: bytes) -> str:
    """Return ``data`` decoded as utf-8 with ANSI escape sequences removed."""

    try:
        text = data.decode("utf-8", errors="replace")
    except Exception:  # pragma: no cover - extremely defensive
        text = data.decode("latin-1", errors="replace")
    # Combine byte-level and string-level strip for safety.
    text = ANSI_RE.sub(b"", data).decode("utf-8", errors="replace")
    # Also drop stray CSI sequences that snuck in as unicode.
    text = re.sub(r"\x1b\[[0-9;?]*[A-Za-z]", "", text)
    return text


# A trivial lock so we don't fork-bomb a remote host if 50 requests land at
# the same time. The lock is per-process; good enough for a single-tenant UI.
_CONCURRENCY = threading.Semaphore(int(os.environ.get("DMS_UI_SSH_CONCURRENCY", "4")))


class DMSSSH:
    """Run commands on a single DMS host.

    Parameters
    ----------
    host_alias:
        The SSH alias as configured in ``~/.ssh/config`` (for example
        ``smail-server``).
    password:
        Optional password override. When ``None`` we fall back to the
        ``SSH_ASKPASS`` script (``/home/hermes/.ssh-askpass.sh``) by setting
        ``SSH_ASKPASS_REQUIRE=force``.
    timeout:
        Per-command timeout in seconds. Default 60s; log tail streams override
        this with their own deadline.
    """

    def __init__(self, host_alias: str, password: Optional[str] = None, timeout: int = 60):
        if not host_alias:
            raise ValueError("host_alias is required")
        self.host_alias = host_alias
        self.password = password
        self.timeout = timeout

    # ---- public API -------------------------------------------------------
    def run(self, command: str, timeout: Optional[int] = None) -> CommandResult:
        """Run ``command`` on the remote host and return its output."""

        deadline = timeout if timeout is not None else self.timeout
        with _CONCURRENCY:
            t0 = time.monotonic()
            if BACKEND == "paramiko":
                result = self._run_paramiko(command, deadline)
            else:
                result = self._run_subprocess(command, deadline)
            log.info(
                "ssh[%s] %r -> rc=%d in %.2fs",
                self.host_alias,
                command[:80],
                result.exit_code,
                time.monotonic() - t0,
            )
            return result

    # ---- paramiko backend -------------------------------------------------
    def _run_paramiko(self, command: str, timeout: int) -> CommandResult:
        assert paramiko is not None
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())  # noqa: S507
        try:
            connect_kwargs: dict = {
                "hostname": self.host_alias,
                "timeout": min(15, timeout),
                "allow_agent": True,
                "look_for_keys": True,
            }
            if self.password:
                connect_kwargs["password"] = self.password
            client.connect(**connect_kwargs)
            stdin, stdout, stderr = client.exec_command(command, timeout=timeout)
            out_b = stdout.read()  # blocks until command exits
            err_b = stderr.read()
            rc = stdout.channel.recv_exit_status()
        except Exception as exc:  # paramiko raises a grab-bag of exceptions
            log.exception("paramiko run failed: %s", exc)
            return CommandResult(stdout="", stderr=str(exc), exit_code=255)
        finally:
            try:
                client.close()
            except Exception:  # pragma: no cover
                pass
        return CommandResult(
            stdout=_strip_ansi(out_b),
            stderr=_strip_ansi(err_b),
            exit_code=rc,
            backend="paramiko",
        )

    # ---- subprocess backend ----------------------------------------------
    def _run_subprocess(self, command: str, timeout: int) -> CommandResult:
        if not ASKPASS_PATH.exists():
            return CommandResult(
                stdout="",
                stderr=f"askpass script missing: {ASKPASS_PATH}",
                exit_code=254,
            )
        env = os.environ.copy()
        env["SSH_ASKPASS"] = str(ASKPASS_PATH)
        env["SSH_ASKPASS_REQUIRE"] = "force"
        env["DISPLAY"] = env.get("DISPLAY", ":0")
        # BatchMode=yes would *suppress* the askpass helper; we want it to be
        # invoked, so we deliberately omit it.
        ssh_argv = [
            "ssh",
            "-T",  # no pty: avoid interactive prompt side-effects
            "-F", str(SSH_CONFIG_PATH),  # use the hermes ssh config so aliases resolve
            "-o",
            "StrictHostKeyChecking=accept-new",
            "-o",
            "UserKnownHostsFile=/dev/null",
            "-o",
            f"ConnectTimeout={min(15, timeout)}",
            self.host_alias,
            command,
        ]
        if self.password:
            env["SSHPASS"] = self.password
            ssh_argv = ["sshpass", "-e"] + ssh_argv  # type: ignore[list-item]
        try:
            completed = subprocess.run(  # noqa: S603 - inputs controlled
                ssh_argv,
                capture_output=True,
                timeout=timeout,
                env=env,
                check=False,
            )
        except subprocess.TimeoutExpired as exc:
            return CommandResult(
                stdout=_strip_ansi(exc.stdout or b""),
                stderr=(exc.stderr or b"").decode("utf-8", errors="replace") + "\n[timeout]",
                exit_code=124,
            )
        except FileNotFoundError as exc:
            return CommandResult(
                stdout="",
                stderr=f"ssh binary not found: {exc}",
                exit_code=127,
            )
        except Exception as exc:  # pragma: no cover - defensive
            return CommandResult(stdout="", stderr=str(exc), exit_code=255)
        return CommandResult(
            stdout=_strip_ansi(completed.stdout),
            stderr=_strip_ansi(completed.stderr),
            exit_code=completed.returncode,
            backend="subprocess",
        )


# ---- streaming helpers used by the SSE log endpoint ----------------------
class StreamingSSH:
    """Open an SSH channel and yield stdout lines until the channel closes.

    Used by the log tail endpoint. The caller (FastAPI) iterates the lines and
    forwards them as Server-Sent Events. Cancellation is cooperative: closing
    the iterator terminates the underlying channel.
    """

    def __init__(self, host_alias: str, command: str, password: Optional[str] = None):
        self.host_alias = host_alias
        self.command = command
        self.password = password
        self._client: Optional["paramiko.SSHClient"] = None  # type: ignore[name-defined]
        self._channel = None

    def __iter__(self) -> "StreamingSSH":
        if BACKEND != "paramiko" or paramiko is None:
            raise RuntimeError(
                "streaming tail requires paramiko backend; "
                f"current backend={BACKEND} import_error={_PARAMIKO_IMPORT_ERROR}"
            )
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())  # noqa: S507
        kwargs = {
            "hostname": self.host_alias,
            "timeout": 15,
            "allow_agent": True,
            "look_for_keys": True,
        }
        if self.password:
            kwargs["password"] = self.password
        client.connect(**kwargs)
        transport = client.get_transport()
        assert transport is not None
        channel = transport.open_session()
        channel.settimeout(1.0)
        channel.exec_command(self.command)
        self._client = client
        self._channel = channel
        return self

    def __next__(self) -> str:
        assert self._channel is not None
        # recv_ready + read so we can periodically check for close + cancellation.
        if self._channel.closed or self._channel.exit_status_ready() and not self._channel.recv_ready():
            raise StopIteration
        try:
            data = self._channel.recv(4096)
        except Exception:
            raise StopIteration
        if not data:
            raise StopIteration
        return _strip_ansi(data)

    def close(self) -> None:
        try:
            if self._channel is not None:
                self._channel.close()
        except Exception:  # pragma: no cover
            pass
        try:
            if self._client is not None:
                self._client.close()
        except Exception:  # pragma: no cover
            pass


def run_command(host_alias: str, command: str, password: Optional[str] = None) -> CommandResult:
    """Functional helper used throughout the FastAPI layer."""

    return DMSSSH(host_alias=host_alias, password=password).run(command)


def quote_args(args: List[str]) -> str:
    """POSIX-shell-quote a list of arguments; safe for ``docker exec`` commands."""

    return " ".join(shlex.quote(a) for a in args)
