#!/usr/bin/env python3
"""SSH MCP server (JSON-RPC 2.0 over stdio).

Exposes SSH tools that operate on remote hosts with **the same output format**
as Skald's native filesystem tools (`read_file`, `list_files`, `grep_files`,
`edit_file`, `replace_lines`, `exec`). The only thing the LLM sees differently
is the first `alias` argument selecting the host. Tool names here are bare
(`read_file`, `exec`, …); Skald prepends the `mcp__ssh__` prefix automatically.

Hosts are addressed by alias — hostname and credentials never appear in tool
calls. Aliases live in ``~/.ssh_aliases.json`` (auto-managed, never edited
by hand). No secret is ever stored in that file.

Login auth (``auth`` per alias, set on ``add_alias``):
  * ``key``      — SSH key / ssh-agent only (default). If the chosen private key
    is encrypted, its passphrase is requested on demand via **MCP elicitation**
    (lazy: only when paramiko reports the key needs one). ``SSH_MCP_KEY_PASSPHRASE``
    still works as a non-interactive override.
  * ``password`` — login password requested on demand via **MCP elicitation**
    (Skald shows a masked field in the Agent Inbox); agent/key auth is skipped.

Elicited login secrets are kept only in this process's RAM with a short TTL
(``SSH_MCP_LOGIN_PW_TTL``), never sent to the LLM and never written to disk;
they are dropped on an authentication failure so the next attempt re-prompts.

sudo (two methods per alias, set on ``add_alias``):
  * ``nopasswd`` — ``sudo -n`` only: non-interactive, fails fast with an
    explicit message if NOPASSWD is not configured on the host (no hung
    channel). No secret stored anywhere.
  * ``prompt``  — ``sudo -n`` is **still tried first**; only if the host really
    demands a password is one requested via **MCP elicitation** (Skald shows a
    masked field in the Agent Inbox), fed to sudo's stdin, kept only in this
    process's RAM with a short TTL, never sent to the LLM and never written to
    disk. ``SSH_MCP_SUDO_PASSWORD`` is a non-interactive override for
    unattended runs (no human to answer the Inbox prompt).

``exec``/``systemd`` never nest sudo: a leading ``sudo`` (with its usual flags,
including ``-u USER``) is stripped from ``command`` and turned into
``sudo=true``, so an agent that writes ``sudo systemctl restart x`` gets the
same, working behaviour as ``sudo=true`` + ``systemctl restart x``. Under sudo
the command runs as ``sh -c '<command>'`` so pipes and redirections are also
privileged.

Connections are pooled per alias with lazy TTL eviction. Host keys are verified
against ``~/.ssh/known_hosts`` (unknown hosts are rejected unless the alias was
added with ``accept_new_host_key=true``).

Run with:
  python3 scripts/ssh_mcp_server.py

Dependency: paramiko>=3.4 (in requirements.txt; installed into .venv by run.sh).
"""

from __future__ import annotations

import itertools
import json
import os
import posixpath
import re
import shlex
import socket
import stat
import sys
import time
from typing import Any


# ── Config ───────────────────────────────────────────────────────────────────

ALIASES_FILE = os.path.expanduser("~/.ssh_aliases.json")

POOL_TTL = int(os.environ.get("SSH_MCP_POOL_TTL", "300"))            # idle connection eviction
SUDO_PW_TTL = int(os.environ.get("SSH_MCP_SUDO_PW_TTL", "300"))      # in-RAM sudo password cache
LOGIN_PW_TTL = int(os.environ.get("SSH_MCP_LOGIN_PW_TTL", "300"))    # in-RAM login/passphrase cache
CONNECT_TIMEOUT = int(os.environ.get("SSH_MCP_CONNECT_TIMEOUT", "15"))
DEFAULT_CMD_TIMEOUT = int(os.environ.get("SSH_MCP_COMMAND_TIMEOUT", "120"))

# Non-interactive sudo password (unattended runs, where nobody can answer the
# elicitation prompt in the Agent Inbox). Empty/unset ⇒ elicitation only.
SUDO_PASSWORD_ENV = os.environ.get("SSH_MCP_SUDO_PASSWORD") or None

# Mirror the native list_files skip set so remote listings match local ones.
SKIP_DIRS = {"target", ".git", "node_modules", ".venv", "__pycache__", "secrets"}

# Match the native read_file cap.
MAX_READ_LINES = 2000


def log(msg: str) -> None:
    """Log to stderr; stdout is reserved for JSON-RPC."""
    print(f"[ssh_mcp] {msg}", file=sys.stderr, flush=True)


class ToolError(Exception):
    """Expected, user-facing failure. Surfaced as ``Error: <message>``."""


# ── stdio JSON-RPC I/O (single readline path so elicit() can re-enter) ─────────

def send(obj: dict) -> None:
    sys.stdout.write(json.dumps(obj) + "\n")
    sys.stdout.flush()


def readline() -> dict | None:
    """Blocking read of one non-empty JSON-RPC message; None on EOF."""
    while True:
        line = sys.stdin.readline()
        if not line:
            return None
        line = line.strip()
        if not line:
            continue
        try:
            return json.loads(line)
        except json.JSONDecodeError as e:
            log(f"invalid JSON input: {e}")
            continue


_eid = itertools.count(1)

# Messages received while blocked on an elicitation reply — replayed by the main
# loop instead of being dropped, so a concurrent tools/call is never lost.
_deferred: list[dict] = []

# Whether the client advertised the `elicitation` capability at initialize. A
# client without it can never supply a password, so we fail fast with a useful
# message instead of blocking on a request nobody will answer.
_client_can_elicit = True


def next_message() -> dict | None:
    """Next inbound message: deferred ones first, then stdin."""
    if _deferred:
        return _deferred.pop(0)
    return readline()


def elicit(message: str, requested_schema: dict) -> dict:
    """Send an ``elicitation/create`` request and block until the reply arrives.

    Returns the JSON-RPC ``result`` ({"action": ..., "content": {...}}). Other
    inbound messages are queued in ``_deferred`` and handled once the reply
    lands (v1: still serial, but nothing is discarded).
    """
    if not _client_can_elicit:
        return {"action": "cancel", "_reason": "unsupported"}
    eid = f"ssh-elicit-{next(_eid)}"
    send({
        "jsonrpc": "2.0",
        "id": eid,
        "method": "elicitation/create",
        "params": {"message": message, "requestedSchema": requested_schema},
    })
    while True:
        msg = readline()
        if msg is None:
            return {"action": "cancel", "_reason": "disconnected"}
        if msg.get("id") == eid:
            return msg.get("result", {"action": "cancel"})
        log(f"deferring inbound while awaiting elicitation: {msg.get('method') or msg.get('id')}")
        _deferred.append(msg)


def _ok(req_id: Any, result: Any) -> dict:
    return {"jsonrpc": "2.0", "id": req_id, "result": result}


def _text_result(req_id: Any, text: str, is_error: bool = False) -> dict:
    res: dict = {"content": [{"type": "text", "text": text}]}
    if is_error:
        res["isError"] = True
    return {"jsonrpc": "2.0", "id": req_id, "result": res}


def _int(args: dict, key: str, default: int) -> int:
    """Integer argument tolerant of null / numeric strings (LLMs send both)."""
    v = args.get(key)
    if v is None or v == "":
        return default
    try:
        return int(v)
    except (TypeError, ValueError):
        return default


# ── Alias store (auto-managed, 0600) ───────────────────────────────────────────

def _load_aliases() -> dict:
    try:
        with open(ALIASES_FILE) as f:
            return json.load(f)
    except FileNotFoundError:
        return {"aliases": []}
    except Exception as e:
        log(f"failed to read aliases: {e}")
        return {"aliases": []}


def _save_aliases(data: dict) -> None:
    os.makedirs(os.path.dirname(ALIASES_FILE), exist_ok=True)
    tmp = f"{ALIASES_FILE}.tmp.{os.getpid()}"
    with open(tmp, "w") as f:
        json.dump(data, f, indent=2)
    os.replace(tmp, ALIASES_FILE)
    try:
        os.chmod(ALIASES_FILE, 0o600)
    except OSError:
        pass


def _find_alias(name: str) -> dict | None:
    for a in _load_aliases().get("aliases", []):
        if a.get("alias") == name:
            return a
    return None


# ── Connection pool (paramiko) ─────────────────────────────────────────────────

_pool: dict[str, dict] = {}            # alias -> {client, sftp, last_used}
_sudo_pw_cache: dict[str, tuple] = {}  # alias -> (password, ts)
_login_pw_cache: dict[str, tuple] = {} # "alias:login" | "alias:passphrase" -> (secret, ts)


def _login_password(alias: str, kind: str = "login") -> str | None:
    """Return the SSH login password (``kind="login"``) or private-key passphrase
    (``kind="passphrase"``) for ``alias`` from the RAM cache, or elicit it.

    Never persisted. Returns None if the user declines/cancels/times out.
    """
    now = time.time()
    key = f"{alias}:{kind}"
    cached = _login_pw_cache.get(key)
    if cached and (now - cached[1] <= LOGIN_PW_TTL):
        return cached[0]

    if kind == "passphrase":
        message = f"Enter the passphrase for the private key of SSH alias '{alias}'."
        title = f"key passphrase — {alias}"
    else:
        message = f"Enter the SSH login password for alias '{alias}'."
        title = f"SSH password — {alias}"

    result = elicit(
        message,
        {
            "type": "object",
            "properties": {
                "password": {"type": "string", "format": "password", "title": title}
            },
            "required": ["password"],
        },
    )
    if result.get("action") == "accept":
        pw = (result.get("content") or {}).get("password", "")
        _login_pw_cache[key] = (pw, now)
        return pw
    return None


def _no_secret_reason() -> str:
    """Why an elicited secret never arrived — the two cases look identical to the
    caller but need very different fixes from the user."""
    return ("this MCP client does not support elicitation" if not _client_can_elicit
            else "the user declined it or the request timed out")


def _clear_login_pw(alias: str) -> None:
    """Drop any cached login password / passphrase for ``alias``."""
    for k in [k for k in _login_pw_cache if k.startswith(f"{alias}:")]:
        _login_pw_cache.pop(k, None)


def _is_auth_failure(paramiko, e: Exception) -> bool:
    """True if ``e`` is an SSH auth rejection a login password could resolve.

    ``AuthenticationException`` (wrong/refused key) always qualifies. A plain
    ``SSHException`` qualifies only when its message says paramiko had no method
    to try — e.g. a password-only host with no key/agent: *"No authentication
    methods available"*. Other SSH errors (banner, host key, protocol) do not.
    """
    if isinstance(e, paramiko.AuthenticationException):
        return True
    msg = str(e).lower()
    return "authentication method" in msg or "no authentication" in msg


def _require_paramiko():
    try:
        import paramiko  # type: ignore
        return paramiko
    except ImportError:
        raise ToolError(
            "paramiko not installed — add 'paramiko>=3.4' to requirements.txt "
            "and reinstall the .venv (uv pip install -r requirements.txt)."
        )


def _connect(cfg: dict, paramiko):
    alias = cfg.get("alias", "")
    auth = (cfg.get("auth") or "key").lower()

    identity = cfg.get("identity_file")
    identity = os.path.expanduser(identity) if identity else None

    password = None
    if auth == "password":
        password = _login_password(alias, "login")
        if password is None:
            raise ToolError(
                f"login password required for alias '{alias}' but none was provided "
                f"({_no_secret_reason()})"
            )

    def attempt(passphrase):
        # With a password in hand, skip agent/key probing so paramiko goes
        # straight to password auth instead of failing on keys first.
        use_pw = password is not None
        client = paramiko.SSHClient()
        client.load_system_host_keys()
        known = os.path.expanduser("~/.ssh/known_hosts")
        if os.path.exists(known):
            try:
                client.load_host_keys(known)
            except Exception as e:
                log(f"could not load known_hosts: {e}")
        if cfg.get("accept_new_host_key"):
            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        else:
            client.set_missing_host_key_policy(paramiko.RejectPolicy())
        client.connect(
            hostname=cfg["hostname"],
            port=int(cfg.get("port", 22)),
            username=cfg.get("username"),
            password=password,
            key_filename=identity,
            passphrase=passphrase,
            allow_agent=not use_pw,
            look_for_keys=not use_pw,
            timeout=CONNECT_TIMEOUT,
        )
        return client

    passphrase = os.environ.get("SSH_MCP_KEY_PASSPHRASE") or None
    try:
        return attempt(passphrase)
    except paramiko.PasswordRequiredException:
        # Encrypted private key with no passphrase supplied — ask for it (lazy).
        if passphrase is not None:
            raise   # we already had one and it was rejected; don't loop
        passphrase = _login_password(alias, "passphrase")
        if passphrase is None:
            raise ToolError(
                f"key passphrase required for alias '{alias}' but none was provided "
                f"({_no_secret_reason()}) — set SSH_MCP_KEY_PASSPHRASE for unattended runs"
            )
        return attempt(passphrase)
    except (paramiko.AuthenticationException, paramiko.SSHException) as e:
        # Key/agent auth was rejected, or the host offers no method paramiko
        # could try (e.g. a password-only host: "No authentication methods
        # available"). If we haven't tried a password yet, elicit one and retry.
        # Declining re-raises the original error. Covers aliases left as the
        # default auth=key that actually need a login password.
        if password is not None or not _is_auth_failure(paramiko, e):
            raise
        password = _login_password(alias, "login")
        if password is None:
            raise
        return attempt(passphrase)


def _close(alias: str) -> None:
    entry = _pool.pop(alias, None)
    if not entry:
        return
    try:
        if entry.get("sftp"):
            entry["sftp"].close()
    except Exception:
        pass
    try:
        entry["client"].close()
    except Exception:
        pass


def _get_client(alias: str):
    cfg = _find_alias(alias)
    if not cfg:
        raise ToolError(f"unknown alias '{alias}'")
    paramiko = _require_paramiko()
    now = time.time()

    entry = _pool.get(alias)
    if entry:
        t = entry["client"].get_transport()
        if (now - entry["last_used"] <= POOL_TTL) and t is not None and t.is_active():
            entry["last_used"] = now
            return entry["client"]
        _close(alias)

    try:
        client = _connect(cfg, paramiko)
    except paramiko.AuthenticationException:
        _clear_login_pw(alias)   # wrong password/passphrase → re-prompt next time
        raise ToolError(f"authentication failed for alias '{alias}' (check key/agent/password)")
    except paramiko.BadHostKeyException:
        raise ToolError(
            f"host key mismatch for alias '{alias}' (possible MITM) — fix ~/.ssh/known_hosts"
        )
    except paramiko.SSHException as e:
        if "not found in known_hosts" in str(e):
            raise ToolError(
                f"unknown host key for alias '{alias}' — re-add it with "
                f"accept_new_host_key=true to trust it on first connect"
            )
        raise ToolError(f"SSH error for alias '{alias}': {e}")
    except (OSError, socket.error) as e:
        raise ToolError(f"connection to alias '{alias}' failed: {e}")

    _pool[alias] = {"client": client, "sftp": None, "last_used": now}
    return client


def _get_sftp(alias: str):
    client = _get_client(alias)
    entry = _pool[alias]
    if entry.get("sftp") is None:
        entry["sftp"] = client.open_sftp()
    return entry["sftp"]


def _run_with_stdin(client, command: str, timeout: int, stdin_data: str | None = None):
    """Run a remote command; return (stdout, stderr, exit_code).

    stdout and stderr are drained **together**: they share one SSH channel
    window, so reading stdout to EOF first stalls as soon as a chatty stderr
    fills that window. ``timeout`` is a wall-clock deadline for the whole run
    (not an idle timeout), and stdin is always closed so commands that read it
    see EOF instead of hanging.
    """
    chan = client.get_transport().open_session(timeout=CONNECT_TIMEOUT)
    try:
        chan.settimeout(timeout)
        chan.exec_command(command)
        try:
            if stdin_data:
                chan.sendall(stdin_data.encode())
            chan.shutdown_write()
        except Exception as e:                    # closed early by the remote end
            log(f"stdin write failed: {e}")

        out, err = bytearray(), bytearray()
        deadline = time.time() + timeout
        settled = 0
        while True:
            idle = True
            while chan.recv_ready():
                out += chan.recv(65536)
                idle = False
            while chan.recv_stderr_ready():
                err += chan.recv_stderr(65536)
                idle = False
            if not idle:
                settled = 0
                continue
            if chan.exit_status_ready():
                # Exit status can arrive before the last data. Wait for the
                # remote EOF, or — should it close without one — for three
                # consecutive empty polls, rather than truncating the output.
                if chan.eof_received or chan.closed or settled >= 3:
                    break
                settled += 1
            if time.time() > deadline:
                raise ToolError(f"command timed out after {timeout}s")
            time.sleep(0.02)
        return (out.decode("utf-8", "replace"),
                err.decode("utf-8", "replace"),
                chan.recv_exit_status())
    except socket.timeout:
        raise ToolError(f"command timed out after {timeout}s")
    finally:
        try:
            chan.close()
        except Exception:
            pass


# ── sudo ───────────────────────────────────────────────────────────────────────

def _cached_sudo_password(alias: str) -> str | None:
    """Live RAM-cache entry for ``alias``, if any.

    Deliberately ignores ``SSH_MCP_SUDO_PASSWORD``: a cache hit lets ``_run_sudo``
    skip its ``sudo -n`` probe, and on a NOPASSWD host sudo would then not read
    the password line at all — feeding it straight into the command's stdin. A
    RAM entry only ever exists because sudo already demanded a password once.
    """
    cached = _sudo_pw_cache.get(alias)
    if cached and (time.time() - cached[1] <= SUDO_PW_TTL):
        return cached[0]
    return None


def _sudo_password(alias: str) -> str | None:
    """Return the sudo password for ``alias`` from RAM cache / env, or elicit it.

    Never persisted. Returns None if the user declines/cancels/times out, or if
    the client cannot elicit at all.
    """
    now = time.time()
    cached = _cached_sudo_password(alias)
    if cached is not None:
        return cached
    if SUDO_PASSWORD_ENV:
        return SUDO_PASSWORD_ENV

    result = elicit(
        f"Enter the sudo password for SSH alias '{alias}'.",
        {
            "type": "object",
            "properties": {
                "password": {
                    "type": "string",
                    "format": "password",
                    "title": f"sudo password — {alias}",
                }
            },
            "required": ["password"],
        },
    )
    if result.get("action") == "accept":
        pw = (result.get("content") or {}).get("password", "")
        _sudo_pw_cache[alias] = (pw, now)
        return pw
    return None


# sudo refusing to run because it wants a password. sudo prints these *instead*
# of running the command, so probing with `-n` is always side-effect free.
_SUDO_NEEDS_PW = (
    "a password is required", "no password was provided", "a terminal is required",
    "no tty present", "askpass",
)
# sudo ran but the password we fed it was wrong.
_SUDO_BAD_PW = ("sorry, try again", "incorrect password attempt", "authentication failure")


def _sudo_says(err: str, code: int, markers: tuple) -> bool:
    """True if a `sudo:`-prefixed stderr line matches one of ``markers``."""
    if code == 0:
        return False
    for line in err.splitlines():
        low = line.strip().lower()
        if low.startswith("sudo:") and any(m in low for m in markers):
            return True
    return False


# Leading `sudo` (plus its common flags) that an agent put in `command` itself.
# Anything unrecognised simply doesn't match and is left untouched.
_SUDO_OPT = (
    r"""(?:-p\s*(?:'[^']*'|"[^"]*"|\S+)"""       # -p PROMPT (takes an argument)
    r"|--prompt(?:=|\s+)\S+"
    r"|-[EHnSbik]+"                              # flag bundles without arguments
    r"|--(?:preserve-env|set-home|non-interactive|stdin|login|shell|background|remove-timestamp))"
)
_LEADING_SUDO = re.compile(
    rf"""^\s*(?:/usr/bin/|/bin/)?sudo
         (?:\s+{_SUDO_OPT})*
         (?:\s+(?:-u\s*|--user(?:=|\s+))(?P<user>[A-Za-z0-9_.\-]+))?
         (?:\s+{_SUDO_OPT})*
         (?:\s+--)?
         \s+(?P<rest>\S.*)$""",
    re.X | re.S,
)


def _split_leading_sudo(command: str) -> tuple[str, bool, str | None]:
    """Split a leading ``sudo …`` off ``command``.

    Returns ``(command_without_sudo, had_sudo, sudo_user)``. Agents routinely
    write ``sudo systemctl restart x``; running that through the sudo machinery
    would nest a second sudo, whose password prompt has no tty and dies. So the
    prefix is peeled off here and expressed as ``sudo=true`` instead.
    """
    m = _LEADING_SUDO.match(command)
    if not m:
        return command, False, None
    rest = m.group("rest").strip()
    if rest.startswith("-"):
        # A sudo option we don't know about — stripping here would hand sudo a
        # mangled command line, so leave the whole thing untouched.
        return command, False, None
    return rest, True, m.group("user")


def _sudo_wrap(flags: str, sudo_user: str | None, command: str) -> str:
    """``sudo <flags> [-u user] sh -c '<command>'`` — the whole command line,
    pipes and redirections included, runs with the elevated privileges."""
    u = f"-u {shlex.quote(sudo_user)} " if sudo_user else ""
    return f"sudo {flags} {u}sh -c {shlex.quote(command)}"


def _run_sudo(alias: str, cfg: dict, client, command: str,
              sudo_user: str | None, timeout: int):
    """Run ``command`` under sudo. Returns (stdout, stderr, exit_code).

    ``sudo -n`` is always attempted first: on a host that grants this user
    NOPASSWD it succeeds outright, so no password is ever requested — which is
    the only thing that works in unattended runs, where nobody is watching the
    Agent Inbox. A password is elicited only when the host actually demands one.
    """
    method = (cfg.get("sudo") or {}).get("method", "prompt")
    if method == "none":
        raise ToolError(
            f"sudo is disabled for alias '{alias}' — re-add it with sudo='prompt' to enable it"
        )

    pw = _cached_sudo_password(alias) if method != "nopasswd" else None
    if pw is None:
        out, err, code = _run_with_stdin(
            client, _sudo_wrap("-n", sudo_user, command), timeout)
        if not _sudo_says(err, code, _SUDO_NEEDS_PW):
            return out, err, code
        if method == "nopasswd":
            raise ToolError(
                f"sudo on '{alias}' requires a password, but the alias is configured with "
                f"sudo='nopasswd' (which only ever runs 'sudo -n' and never prompts). "
                f"Re-add the alias with sudo='prompt', or grant this user a NOPASSWD rule "
                f"in the remote /etc/sudoers."
            )
        pw = _sudo_password(alias)
        if pw is None:
            raise ToolError(
                f"sudo password required for alias '{alias}' but none was provided "
                f"({_no_secret_reason()}). Either answer the sudo prompt in the Agent Inbox, or — for "
                f"unattended runs — set SSH_MCP_SUDO_PASSWORD in the connector settings, or "
                f"grant this user NOPASSWD sudo on the host."
            )

    out, err, code = _run_with_stdin(
        client, _sudo_wrap("-S -p ''", sudo_user, command), timeout, pw + "\n")
    if _sudo_says(err, code, _SUDO_BAD_PW) or _sudo_says(err, code, _SUDO_NEEDS_PW):
        _sudo_pw_cache.pop(alias, None)     # drop it so the next call re-prompts
        raise ToolError(
            f"sudo password rejected on '{alias}'"
            + (" (from SSH_MCP_SUDO_PASSWORD)" if SUDO_PASSWORD_ENV else
               " — the cached password was discarded, retry to be asked again")
        )
    return out, err, code


# ── SFTP helpers ───────────────────────────────────────────────────────────────

def _sftp_read_text(sftp, path: str) -> str:
    with sftp.open(path, "r") as f:
        data = f.read()
    return data.decode("utf-8", "replace") if isinstance(data, (bytes, bytearray)) else data


def _sftp_write_atomic(sftp, path: str, content: str) -> None:
    """Write atomically: temp file in the same dir + posix_rename. Preserve mode."""
    d = posixpath.dirname(path) or "."
    base = posixpath.basename(path)
    tmp = posixpath.join(d, f".{base}.tmp.{os.getpid()}")

    mode = None
    try:
        mode = stat.S_IMODE(sftp.stat(path).st_mode)
    except IOError:
        pass

    with sftp.open(tmp, "w") as f:
        f.write(content)
    if mode is not None:
        try:
            sftp.chmod(tmp, mode)
        except IOError:
            pass
    try:
        sftp.posix_rename(tmp, path)
    except (IOError, AttributeError):
        try:
            sftp.remove(path)
        except IOError:
            pass
        sftp.rename(tmp, path)


def _sftp_mkdirs(sftp, d: str) -> None:
    if not d or d in ("/", "."):
        return
    try:
        sftp.stat(d)
        return
    except IOError:
        pass
    parent = posixpath.dirname(d)
    if parent and parent != d:
        _sftp_mkdirs(sftp, parent)
    try:
        sftp.mkdir(d)
    except IOError:
        pass


def _relpath(root: str, full: str) -> str:
    r = root.rstrip("/") or "/"
    return posixpath.relpath(full, r)


# ── Tools: aliases ─────────────────────────────────────────────────────────────

def _tool_list_aliases(args: dict) -> str:
    out = []
    for a in _load_aliases().get("aliases", []):
        out.append({
            "alias": a.get("alias"),
            "hostname": a.get("hostname"),
            "port": a.get("port", 22),
            "username": a.get("username"),
            "auth": a.get("auth", "key"),
            "sudo_method": (a.get("sudo") or {}).get("method", "prompt"),
            "description": a.get("description", ""),
        })
    return json.dumps(out, indent=2)


def _tool_add_alias(args: dict) -> str:
    name = args.get("alias")
    if not name:
        return "Error: missing required argument: alias"
    if not args.get("hostname"):
        return "Error: missing required argument: hostname"

    sudo = args.get("sudo")
    method = sudo.get("method") if isinstance(sudo, dict) else (sudo or "prompt")
    if method not in ("nopasswd", "prompt", "none"):
        return f"Error: invalid sudo method '{method}' (use nopasswd|prompt|none)"

    auth = (args.get("auth") or "key").lower()
    if auth not in ("key", "password"):
        return f"Error: invalid auth method '{auth}' (use key|password)"

    entry = {
        "alias": name,
        "hostname": args["hostname"],
        "port": int(args.get("port", 22)),
        "username": args.get("username"),
        "identity_file": args.get("identity_file"),
        "description": args.get("description", ""),
        "auth": auth,
        "sudo": {"method": method},
        "accept_new_host_key": bool(args.get("accept_new_host_key", False)),
    }

    data = _load_aliases()
    aliases = data.setdefault("aliases", [])
    prev = None
    for i, a in enumerate(aliases):
        if a.get("alias") == name:
            prev = a
            aliases[i] = entry
            break
    else:
        aliases.append(entry)
    _save_aliases(data)
    _close(name)          # config may have changed — drop any pooled connection
    _sudo_pw_cache.pop(name, None)
    _clear_login_pw(name)

    target = f"{entry.get('username')}@{entry['hostname']}:{entry['port']}"
    if prev:
        return f"Updated alias '{name}' → {target} (auth: {auth}, sudo: {method})."
    return f"Added alias '{name}' → {target} (auth: {auth}, sudo: {method})."


def _tool_remove_alias(args: dict) -> str:
    name = args.get("alias")
    if not name:
        return "Error: missing required argument: alias"
    data = _load_aliases()
    aliases = data.get("aliases", [])
    kept = [a for a in aliases if a.get("alias") != name]
    if len(kept) == len(aliases):
        return f"Error: alias '{name}' not found"
    data["aliases"] = kept
    _save_aliases(data)
    _close(name)
    _sudo_pw_cache.pop(name, None)
    _clear_login_pw(name)
    return f"Removed alias '{name}'."


# ── Tools: filesystem (native output format) ───────────────────────────────────

def _tool_read_file(args: dict) -> str:
    alias, path = args.get("alias"), args.get("path")
    if not alias or not path:
        return "Error: 'alias' and 'path' are required"
    sftp = _get_sftp(alias)
    try:
        content = _sftp_read_text(sftp, path)
    except IOError as e:
        raise ToolError(f"cannot read {path}: {e}")

    lines = content.splitlines()
    total = len(lines)

    limit = args.get("limit")
    limit = min(int(limit), MAX_READ_LINES) if limit is not None else None
    start = max(int(args["start_line"]) - 1, 0) if args.get("start_line") is not None else 0
    if args.get("end_line") is not None:
        end = min(int(args["end_line"]), total)
    elif limit is not None:
        end = min(start + limit, total)
    else:
        end = total

    if start >= total and total > 0:
        return f"(file has only {total} lines; start_line {start + 1} is out of range)"
    end = max(end, start)

    width = max(len(str(total)), 3)
    return "\n".join(
        f"{start + i + 1:>{width}} | {line}" for i, line in enumerate(lines[start:end])
    )


def _tool_list_files(args: dict) -> str:
    alias, path = args.get("alias"), args.get("path")
    if not alias or not path:
        return "Error: 'alias' and 'path' are required"
    max_depth = _int(args, "depth", 3)
    dirs_only = bool(args.get("dirs_only", False))
    sftp = _get_sftp(alias)

    out: list[str] = []

    def walk(d: str, depth: int) -> None:
        try:
            entries = sftp.listdir_attr(d)
        except IOError:
            return
        for a in entries:
            full = posixpath.join(d, a.filename)
            if stat.S_ISDIR(a.st_mode):
                if a.filename in SKIP_DIRS:
                    continue
                if dirs_only:
                    out.append(_relpath(path, full))
                if depth + 1 < max_depth:
                    walk(full, depth + 1)
            elif stat.S_ISREG(a.st_mode) and not dirs_only:
                out.append(_relpath(path, full))

    try:
        sftp.listdir_attr(path)
    except IOError as e:
        raise ToolError(f"cannot list {path}: {e}")
    walk(path, 0)
    out.sort()
    return json.dumps(out)


def _grep_flags(args: dict) -> str:
    flags = ""
    if not bool(args.get("case_sensitive", False)):
        flags += "-i "
    inc = args.get("include_glob")
    if inc:
        flags += f"--include={shlex.quote(inc)} "
    return flags


def _tool_grep_files(args: dict) -> str:
    alias, path, pattern = args.get("alias"), args.get("path"), args.get("pattern")
    if not alias or not path or pattern is None:
        return "Error: 'alias', 'path' and 'pattern' are required"
    mode = args.get("output_mode", "content")
    ctx = min(_int(args, "context_lines", 0), 10)
    maxr = _int(args, "max_results", 100)
    client = _get_client(alias)

    flags = _grep_flags(args)
    qpat, qpath = shlex.quote(pattern), shlex.quote(path)
    root_prefix = path.rstrip("/") + "/"

    def rel(p: str) -> str:
        return p[len(root_prefix):] if p.startswith(root_prefix) else p

    if mode == "files_only":
        cmd = f"grep -rlIZ {flags}-E -e {qpat} -- {qpath}"
        out, err, code = _run_with_stdin(client, cmd, DEFAULT_CMD_TIMEOUT)
        if code >= 2 and not out:
            raise ToolError(err.strip() or "grep failed")
        files = [rel(f) for f in out.split("\0") if f][:maxr]
        if not files:
            return f'No files match "{pattern}" in {path}.'
        return f"{len(files)} file(s):\n" + "\n".join(files)

    if mode == "count":
        cmd = f"grep -rcI {flags}-E -e {qpat} -- {qpath}"
        out, err, code = _run_with_stdin(client, cmd, DEFAULT_CMD_TIMEOUT)
        if code >= 2 and not out:
            raise ToolError(err.strip() or "grep failed")
        items = []
        for line in out.splitlines():
            f, _, c = line.rpartition(":")     # rpartition: count is numeric at end
            if f and c.isdigit() and int(c) > 0:
                items.append((rel(f), int(c)))
        items = items[:maxr]
        if not items:
            return f'No matches for "{pattern}" in {path}.'
        return f"{len(items)} file(s):\n" + "\n".join(f"{f}: {c}" for f, c in items)

    # content mode
    cflag = f"-C {ctx} " if ctx else ""
    cmd = f"grep -rnIZ {cflag}{flags}-E -e {qpat} -- {qpath}"
    out, err, code = _run_with_stdin(client, cmd, DEFAULT_CMD_TIMEOUT)
    if code >= 2 and not out:
        raise ToolError(err.strip() or "grep failed")

    entries: list[str] = []
    if ctx == 0:
        for line in out.split("\n"):
            if not line:
                continue
            if "\0" in line:
                f, _, rest = line.partition("\0")
            else:
                f, _, rest = line.partition(":")
            lineno, _, body = rest.partition(":")
            entries.append(f"{rel(f)}:{lineno}: {body}")
            if len(entries) >= maxr:
                break
    else:
        prev_file = None
        for line in out.split("\n"):
            if not line:
                continue
            if line == "--":
                if prev_file is not None and len(entries) < maxr:
                    entries.append(f"{rel(prev_file)}:---")
                continue
            if "\0" in line:
                f, _, rest = line.partition("\0")
            else:
                f, _, rest = line.partition(":")
            m = re.match(r"(\d+)([:-])(.*)$", rest, re.S)
            if not m:
                continue
            lineno, sep, body = m.group(1), m.group(2), m.group(3)
            marker = ">" if sep == ":" else " "
            entries.append(f"{marker}{rel(f)}: {lineno}: {body}")
            prev_file = f
            if len(entries) >= maxr:
                break

    if not entries:
        return f'No matches for "{pattern}" in {path}.'
    return f"{len(entries)} match(es):\n" + "\n".join(entries)


def _tool_edit_file(args: dict) -> str:
    alias, path = args.get("alias"), args.get("path")
    old, new = args.get("old"), args.get("new")
    if not alias or not path:
        return "Error: 'alias' and 'path' are required"
    if old is None or new is None:
        return "Error: 'old' and 'new' are required"
    replace_all = bool(args.get("replace_all", False))
    sftp = _get_sftp(alias)
    try:
        content = _sftp_read_text(sftp, path)
    except IOError as e:
        raise ToolError(f"cannot read {path}: {e}")

    not_found = (
        f"Error: Text not found in {path}. "
        f"Call read_file first and copy the text exactly as shown after the '| ' prefix."
    )
    if replace_all:
        if old not in content:
            return not_found
        updated = content.replace(old, new)
    else:
        cnt = content.count(old)
        if cnt > 1:
            return (
                f"Error: Text found {cnt} times in {path}. "
                f"Include more surrounding context in `old` to make it unique, "
                f"or set replace_all=true."
            )
        if cnt == 0:
            return not_found
        updated = content.replace(old, new, 1)

    _sftp_write_atomic(sftp, path, updated)
    return f"Edited {path}."


def _tool_replace_lines(args: dict) -> str:
    alias, path = args.get("alias"), args.get("path")
    if not alias or not path:
        return "Error: 'alias' and 'path' are required"
    if args.get("from_line") is None or args.get("to_line") is None or args.get("new") is None:
        return "Error: 'from_line', 'to_line' and 'new' are required"
    from_line = int(args["from_line"])
    to_line = int(args["to_line"])
    new = args["new"]
    if from_line < 1:
        return "Error: from_line must be >= 1"
    if to_line < from_line:
        return "Error: to_line must be >= from_line"

    sftp = _get_sftp(alias)
    try:
        content = _sftp_read_text(sftp, path)
    except IOError as e:
        raise ToolError(f"cannot read {path}: {e}")

    lines = content.splitlines()
    total = len(lines)
    if from_line > total:
        return f"Error: from_line {from_line} exceeds file length ({total} lines)"
    to_clamped = min(to_line, total)
    new_lines = new.splitlines()
    lines[from_line - 1:to_clamped] = new_lines

    updated = "\n".join(lines)
    if content.endswith("\n"):
        updated += "\n"
    _sftp_write_atomic(sftp, path, updated)
    return f"Replaced lines {from_line}–{to_clamped} in {path} with {len(new_lines)} new lines."


# ── Tools: exec / sudo / systemd ────────────────────────────────────────────────

def _tool_exec(args: dict) -> str:
    alias, command = args.get("alias"), args.get("command")
    if not alias or command is None:
        return "Error: 'alias' and 'command' are required"
    sudo = bool(args.get("sudo", False))
    sudo_user = args.get("sudo_user")
    timeout = _int(args, "timeout_sec", DEFAULT_CMD_TIMEOUT)
    cfg = _find_alias(alias)
    if not cfg:
        return f"Error: unknown alias '{alias}'"

    # A `sudo` the agent typed into `command` is the same intent as sudo=true —
    # honour it here rather than nesting a second, tty-less sudo remotely.
    command, inline_sudo, inline_user = _split_leading_sudo(command)
    if inline_sudo:
        sudo = True
        sudo_user = sudo_user or inline_user
    if sudo_user:
        sudo = True          # `sudo -u X` is meaningless without sudo itself
    if not command.strip():
        return "Error: 'command' is empty"

    client = _get_client(alias)
    if sudo:
        out, err, code = _run_sudo(alias, cfg, client, command, sudo_user, timeout)
    else:
        out, err, code = _run_with_stdin(client, command, timeout)

    result = {"stdout": out, "stderr": err, "exit_code": code}
    if _sudo_says(err, code, _SUDO_NEEDS_PW):
        # A sudo buried mid-command (e.g. `cd /x && sudo …`) that we could not
        # peel off. Don't silently re-run the whole line as root — tell the LLM.
        result["hint"] = (
            "this command used sudo internally and sudo asked for a password on a "
            "tty it does not have. Re-run it with sudo=true and no 'sudo' inside "
            "`command` (with sudo=true the whole command line runs as root)."
        )
    return json.dumps(result)


def _tool_systemd(args: dict) -> str:
    alias, service, action = args.get("alias"), args.get("service"), args.get("action")
    if not alias or not service or not action:
        return "Error: 'alias', 'service' and 'action' are required"
    allowed = {"status", "start", "stop", "restart", "reload", "enable", "disable"}
    if action not in allowed:
        return f"Error: invalid action '{action}' (allowed: {', '.join(sorted(allowed))})"
    cfg = _find_alias(alias)
    if not cfg:
        return f"Error: unknown alias '{alias}'"

    qsvc = shlex.quote(service)
    client = _get_client(alias)

    parts: list[str] = []
    if action != "status":
        out, err, code = _run_sudo(
            alias, cfg, client, f"systemctl {action} {qsvc}", None, DEFAULT_CMD_TIMEOUT)
        parts.append(f"$ systemctl {action} {service}  (exit {code})")
        if out.strip():
            parts.append(out.strip())
        if err.strip():
            parts.append(err.strip())

    status, _, _ = _run_with_stdin(
        client, f"systemctl status {qsvc} --no-pager 2>&1 | head -n 20", DEFAULT_CMD_TIMEOUT)
    parts.append("── status ──")
    parts.append(status.strip())

    journal, _, _ = _run_with_stdin(
        client, f"journalctl -u {qsvc} -n 10 --no-pager 2>&1", DEFAULT_CMD_TIMEOUT)
    parts.append("── journal (last 10) ──")
    parts.append(journal.strip())
    return "\n".join(parts)


# ── Tools: transfer / diagnostics ───────────────────────────────────────────────

def _tool_upload(args: dict) -> str:
    alias = args.get("alias")
    local_path, remote_path = args.get("local_path"), args.get("remote_path")
    if not alias or not local_path or not remote_path:
        return "Error: 'alias', 'local_path' and 'remote_path' are required"
    if not os.path.exists(local_path):
        return f"Error: local path not found: {local_path}"
    sftp = _get_sftp(alias)

    count = total = 0
    dest_shown = remote_path
    if os.path.isdir(local_path):
        for root, _dirs, files in os.walk(local_path):
            relroot = os.path.relpath(root, local_path)
            rdir = remote_path if relroot == "." else posixpath.join(
                remote_path, relroot.replace(os.sep, "/"))
            _sftp_mkdirs(sftp, rdir)
            for fn in files:
                lf = os.path.join(root, fn)
                sftp.put(lf, posixpath.join(rdir, fn))
                count += 1
                total += os.path.getsize(lf)
    else:
        # scp/rsync semantics: a trailing-slash or existing-directory remote_path
        # means "upload the file INTO that directory". paramiko's sftp.put needs a
        # full destination FILE path — handed a directory path it fails with a
        # generic "Failure" — so append the local basename in that case.
        into_dir = remote_path.endswith("/")
        if not into_dir:
            try:
                into_dir = stat.S_ISDIR(sftp.stat(remote_path).st_mode)
            except IOError:
                into_dir = False
        if into_dir:
            target_dir = remote_path.rstrip("/") or "/"
            _sftp_mkdirs(sftp, target_dir)
            dest = posixpath.join(target_dir, os.path.basename(local_path))
        else:
            dest = remote_path
            parent = posixpath.dirname(dest)
            if parent:
                _sftp_mkdirs(sftp, parent)
        sftp.put(local_path, dest)
        count, total = 1, os.path.getsize(local_path)
        dest_shown = dest
    return f"Uploaded {count} file(s), {total} bytes → {dest_shown}"


def _tool_download(args: dict) -> str:
    alias = args.get("alias")
    remote_path, local_path = args.get("remote_path"), args.get("local_path")
    if not alias or not remote_path or not local_path:
        return "Error: 'alias', 'remote_path' and 'local_path' are required"
    sftp = _get_sftp(alias)
    try:
        st = sftp.stat(remote_path)
    except IOError as e:
        raise ToolError(f"remote path not found: {remote_path} ({e})")

    count = total = 0
    if stat.S_ISDIR(st.st_mode):
        def rec(rdir: str, ldir: str) -> None:
            nonlocal count, total
            os.makedirs(ldir, exist_ok=True)
            for a in sftp.listdir_attr(rdir):
                rf = posixpath.join(rdir, a.filename)
                lf = os.path.join(ldir, a.filename)
                if stat.S_ISDIR(a.st_mode):
                    rec(rf, lf)
                elif stat.S_ISREG(a.st_mode):
                    sftp.get(rf, lf)
                    count += 1
                    total += a.st_size or os.path.getsize(lf)
        rec(remote_path, local_path)
    else:
        parent = os.path.dirname(local_path)
        if parent:
            os.makedirs(parent, exist_ok=True)
        sftp.get(remote_path, local_path)
        count, total = 1, os.path.getsize(local_path)
    return f"Downloaded {count} file(s), {total} bytes → {local_path}"


def _tool_sysinfo(args: dict) -> str:
    alias = args.get("alias")
    if not alias:
        return "Error: 'alias' is required"
    client = _get_client(alias)
    cmd = (
        "echo OS=$(uname -s 2>/dev/null); "
        "echo KERNEL=$(uname -r 2>/dev/null); "
        "echo CPU=$(nproc 2>/dev/null); "
        "echo MEMTOTAL=$(awk '/MemTotal/{print $2}' /proc/meminfo 2>/dev/null); "
        "echo MEMAVAIL=$(awk '/MemAvailable/{print $2}' /proc/meminfo 2>/dev/null); "
        "echo DISKTOTAL=$(df -kP / 2>/dev/null | tail -1 | awk '{print $2}'); "
        "echo DISKAVAIL=$(df -kP / 2>/dev/null | tail -1 | awk '{print $4}'); "
        "echo UPTIME=$(uptime -p 2>/dev/null || uptime 2>/dev/null)"
    )
    out, _, _ = _run_with_stdin(client, cmd, DEFAULT_CMD_TIMEOUT)
    kv: dict[str, str] = {}
    for line in out.splitlines():
        if "=" in line:
            k, _, v = line.partition("=")
            kv[k.strip()] = v.strip()

    def gb(key: str):
        try:
            return round(int(kv.get(key, "")) / 1024 / 1024, 2)
        except (ValueError, TypeError):
            return None

    info = {
        "os": kv.get("OS", ""),
        "kernel": kv.get("KERNEL", ""),
        "cpu_count": int(kv["CPU"]) if kv.get("CPU", "").isdigit() else None,
        "ram_total_gb": gb("MEMTOTAL"),
        "ram_free_gb": gb("MEMAVAIL"),
        "disk_total_gb": gb("DISKTOTAL"),
        "disk_free_gb": gb("DISKAVAIL"),
        "uptime": kv.get("UPTIME", ""),
    }
    return json.dumps(info, indent=2)


# ── Tool registry ────────────────────────────────────────────────────────────────

_ALIAS = {"type": "string", "description": "Host alias registered via add_alias."}
_SFTP_NOTE = (
    " Runs as the login user (no sudo): for paths needing root, use exec with "
    "sudo=true (e.g. tee/install)."
)

TOOLS = [
    {
        "name": "list_aliases",
        "title": "List Aliases",
        "description": "List configured SSH host aliases (never reveals keys or sudo passwords).",
        "inputSchema": {"type": "object", "properties": {}},
    },
    {
        "name": "add_alias",
        "title": "Add Alias",
        "description": "Register or update an SSH host alias. Login via SSH key/agent (default) or login password asked on demand via elicitation.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "alias": {"type": "string", "description": "Short name used to address the host."},
                "hostname": {"type": "string", "description": "Host or IP."},
                "port": {"type": "integer", "description": "SSH port (default 22)."},
                "username": {"type": "string", "description": "Login user."},
                "identity_file": {"type": "string", "description": "Path to private key (optional; ssh-agent is also tried). An encrypted key's passphrase is asked via elicitation."},
                "description": {"type": "string", "description": "Free-text note."},
                "auth": {"type": "string", "enum": ["key", "password"],
                         "description": "Login auth. key: SSH key/agent (default). password: login password asked on demand via elicitation, kept only in RAM."},
                "sudo": {"type": "string", "enum": ["nopasswd", "prompt", "none"],
                         "description": "How sudo authenticates. Keep the default 'prompt' unless you KNOW otherwise: it tries 'sudo -n' first (so a host with a NOPASSWD rule never prompts) and only asks the user for the sudo password via elicitation if the host actually demands one. 'nopasswd' runs 'sudo -n' only and NEVER prompts — pick it just to forbid prompting, since on a host without a NOPASSWD rule every sudo call then fails. 'none' disables sudo entirely. Default prompt."},
                "accept_new_host_key": {"type": "boolean", "description": "Trust the host key on first connect (TOFU). Default false."},
            },
            "required": ["alias", "hostname", "username"],
        },
    },
    {
        "name": "remove_alias",
        "title": "Remove Alias",
        "description": "Remove a host alias and close its pooled connection.",
        "inputSchema": {"type": "object", "properties": {"alias": _ALIAS}, "required": ["alias"]},
    },
    {
        "name": "read_file",
        "title": "Read File",
        "description": "Read a remote file with 1-based line numbers (same format as the local read_file)." + _SFTP_NOTE,
        "inputSchema": {
            "type": "object",
            "properties": {
                "alias": _ALIAS,
                "path": {"type": "string", "description": "Absolute remote path."},
                "start_line": {"type": "integer", "description": "First line (1-based, inclusive)."},
                "end_line": {"type": "integer", "description": "Last line (1-based, inclusive)."},
                "limit": {"type": "integer", "description": "Max lines to read (cap 2000)."},
            },
            "required": ["alias", "path"],
        },
    },
    {
        "name": "list_files",
        "title": "List Files",
        "description": "List files/dirs under a remote path; returns a JSON array of relative paths (same as local list_files).",
        "inputSchema": {
            "type": "object",
            "properties": {
                "alias": _ALIAS,
                "path": {"type": "string", "description": "Absolute remote directory."},
                "depth": {"type": "integer", "description": "Max recursion depth (default 3; 1 = immediate contents)."},
                "dirs_only": {"type": "boolean", "description": "Only directories (default false)."},
            },
            "required": ["alias", "path"],
        },
    },
    {
        "name": "grep_files",
        "title": "Grep Files",
        "description": "Search a remote path with a regex; output matches the local grep_files (uses remote grep -E).",
        "inputSchema": {
            "type": "object",
            "properties": {
                "alias": _ALIAS,
                "path": {"type": "string", "description": "Remote file or directory."},
                "pattern": {"type": "string", "description": "Regex (case-insensitive by default)."},
                "case_sensitive": {"type": "boolean", "description": "Default false."},
                "include_glob": {"type": "string", "description": "Restrict to files matching this glob, e.g. '*.rs'."},
                "output_mode": {"type": "string", "enum": ["content", "files_only", "count"], "description": "Default 'content'."},
                "context_lines": {"type": "integer", "description": "Lines of context per match (default 0, max 10)."},
                "max_results": {"type": "integer", "description": "Stop after N results (default 100)."},
            },
            "required": ["alias", "path", "pattern"],
        },
    },
    {
        "name": "edit_file",
        "title": "Edit File",
        "description": "Find & replace in a remote file (atomic). `old` must match exactly once unless replace_all." + _SFTP_NOTE,
        "inputSchema": {
            "type": "object",
            "properties": {
                "alias": _ALIAS,
                "path": {"type": "string", "description": "Absolute remote path."},
                "old": {"type": "string", "description": "Exact text to replace."},
                "new": {"type": "string", "description": "Replacement text."},
                "replace_all": {"type": "boolean", "description": "Replace every occurrence (default false)."},
            },
            "required": ["alias", "path", "old", "new"],
        },
    },
    {
        "name": "replace_lines",
        "title": "Replace Lines",
        "description": "Replace a 1-based inclusive line range in a remote file (atomic)." + _SFTP_NOTE,
        "inputSchema": {
            "type": "object",
            "properties": {
                "alias": _ALIAS,
                "path": {"type": "string", "description": "Absolute remote path."},
                "from_line": {"type": "integer", "description": "First line (1-based, inclusive)."},
                "to_line": {"type": "integer", "description": "Last line (1-based, inclusive)."},
                "new": {"type": "string", "description": "Replacement text."},
            },
            "required": ["alias", "path", "from_line", "to_line", "new"],
        },
    },
    {
        "name": "exec",
        "title": "Execute Command",
        "description": (
            "Run a command on the remote host. For anything needing root, set sudo=true and "
            "write the command WITHOUT a 'sudo' prefix — with sudo=true the entire command "
            "line (pipes and redirections included) runs as root. A leading 'sudo' left in "
            "`command` is stripped and treated as sudo=true anyway, but a 'sudo' in the middle "
            "of a command (e.g. 'cd /x && sudo …') has no tty and will fail."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "alias": _ALIAS,
                "command": {"type": "string", "description": "Shell command. Do NOT prefix it with 'sudo' — use the sudo argument."},
                "sudo": {"type": "boolean", "description": "Run the whole command as root, via sudo (default false)."},
                "sudo_user": {"type": "string", "description": "Target user for sudo -u (optional; implies sudo)."},
                "timeout_sec": {"type": "integer", "description": "Kill after N seconds (default 120)."},
            },
            "required": ["alias", "command"],
        },
    },
    {
        "name": "upload",
        "title": "Upload File",
        "description": "Upload a local file or directory (recursive) to the remote host via SFTP." + _SFTP_NOTE,
        "inputSchema": {
            "type": "object",
            "properties": {
                "alias": _ALIAS,
                "local_path": {"type": "string", "description": "Local file or directory."},
                "remote_path": {"type": "string", "description": "Remote destination. For a single file: a trailing '/' (or an existing remote directory) uploads the file INTO that directory keeping its name; otherwise it is the exact destination file path (parent dirs are created)."},
            },
            "required": ["alias", "local_path", "remote_path"],
        },
    },
    {
        "name": "download",
        "title": "Download File",
        "description": "Download a remote file or directory (recursive) to the local host via SFTP.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "alias": _ALIAS,
                "remote_path": {"type": "string", "description": "Remote file or directory."},
                "local_path": {"type": "string", "description": "Local destination path."},
            },
            "required": ["alias", "remote_path", "local_path"],
        },
    },
    {
        "name": "sysinfo",
        "title": "System Info",
        "description": "Report OS, kernel, CPU count, RAM and root-disk usage, and uptime.",
        "inputSchema": {"type": "object", "properties": {"alias": _ALIAS}, "required": ["alias"]},
    },
    {
        "name": "systemd",
        "title": "Manage Systemd Service",
        "description": "Manage a systemd service (status/start/stop/restart/reload/enable/disable) + last 10 journal lines. Mutating actions use sudo.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "alias": _ALIAS,
                "service": {"type": "string", "description": "Service/unit name."},
                "action": {"type": "string", "enum": ["status", "start", "stop", "restart", "reload", "enable", "disable"]},
            },
            "required": ["alias", "service", "action"],
        },
    },
]

TOOL_DISPATCH = {
    "list_aliases": _tool_list_aliases,
    "add_alias": _tool_add_alias,
    "remove_alias": _tool_remove_alias,
    "read_file": _tool_read_file,
    "list_files": _tool_list_files,
    "grep_files": _tool_grep_files,
    "edit_file": _tool_edit_file,
    "replace_lines": _tool_replace_lines,
    "exec": _tool_exec,
    "upload": _tool_upload,
    "download": _tool_download,
    "sysinfo": _tool_sysinfo,
    "systemd": _tool_systemd,
}


# ── JSON-RPC dispatch ────────────────────────────────────────────────────────────

def handle_message(msg: dict) -> dict | None:
    method = msg.get("method", "")
    req_id = msg.get("id")

    if method == "initialize":
        # Remember whether the client can collect input mid-call: without the
        # elicitation capability we must never block on a password prompt.
        global _client_can_elicit
        caps = (msg.get("params") or {}).get("capabilities") or {}
        _client_can_elicit = "elicitation" in caps
        log(f"client elicitation capability: {'yes' if _client_can_elicit else 'no'}")
        return _ok(req_id, {
            "protocolVersion": "2025-06-18",
            "capabilities": {"tools": {}},
            "serverInfo": {"name": "ssh", "version": "1.1.0"},
        })
    if method == "notifications/initialized":
        return None
    if method == "tools/list":
        return _ok(req_id, {"tools": TOOLS})
    if method == "tools/call":
        params = msg.get("params", {})
        name = params.get("name", "")
        targs = params.get("arguments", {}) or {}
        handler = TOOL_DISPATCH.get(name)
        if handler is None:
            return _text_result(req_id, f"Error: Unknown tool: {name}", True)
        try:
            text = handler(targs)
        except ToolError as e:
            text = f"Error: {e}"
        except Exception as e:
            log(f"unhandled exception in tool '{name}': {e}")
            text = f"Error: internal error in '{name}': {e}"
        return _text_result(req_id, text, text.startswith("Error:"))

    if req_id is not None:
        return {"jsonrpc": "2.0", "id": req_id,
                "error": {"code": -32601, "message": f"Method not found: {method}"}}
    return None


def main() -> None:
    log("starting SSH MCP server")
    try:
        while True:
            msg = next_message()
            if msg is None:
                break
            resp = handle_message(msg)
            if resp is not None:
                send(resp)
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()