"""One ssh layer for albedo: the daemon's remote kernels and the model's remote plugin.

Both ride the same multiplexed control connections (`/tmp/albedo-ssh-cm/%C`),
so a host warmed by one is warm for the other. BatchMode stays on: a host
that needs a passphrase, a second factor or a new host key answers
`needs_auth`, and the tui opens the master in the terminal instead.

    albedo_ssh.py probe <[user@]host>
    albedo_ssh.py commands <[user@]host> <home>
    albedo_ssh.py hosts

`probe` prints `{"step": "staging"}` on a line of its own before it copies
the bundle over, so a waiting client can say so, then one JSON object:
`{state, detail, os, arch, home, cpus, python}` and, once the home is known,
the `commands` answer: `{argv, control, bundle, bridge, signal, remove,
gather, auth_sock}`, the complete remote commands the daemon runs there.
state is ready, needs_auth, unreachable or unsupported. A ready probe has checked python >= 3.11 and staged the content-hashed bundle
under `~/.albedo-remote/<digest>`. `commands` needs no network, for a host
known only from a kernel's record while ssh is down. `hosts` lists the `Host`
names of ~/.ssh/config and its includes, patterns left out.
"""

from __future__ import annotations

import glob
import io
import json
import os
import shlex
import subprocess
import sys
import tarfile
import tempfile
from collections.abc import Callable

import albedo_bundle

CONNECT_TIMEOUT = 10  # seconds ssh may take to reach a host
PROBE_TIMEOUT = 30  # seconds for the whole probe command
STAGE_TIMEOUT = 120  # seconds to move the bundle over one ssh stream
MINIMUM_PYTHON = (3, 11)
STAGED_DIR = ".albedo-remote"


def control_dir() -> str:
    """Control sockets live under /tmp when it fits, else the temp dir: unix
    socket paths are capped (~104 bytes on macOS) and %C hashes to 40 more,
    so neither ALBEDO_HOME nor macOS's deep $TMPDIR may be on the path."""
    control = os.path.join("/tmp", "albedo-ssh-cm")
    if len(control) + 42 > 104:
        control = os.path.join(tempfile.gettempdir(), "albedo-ssh-cm")
    return control


def base() -> list[str]:
    """The ssh argv prefix: one multiplexed control connection per target."""
    control = control_dir()
    os.makedirs(control, exist_ok=True)
    return [
        "ssh",
        "-o",
        "BatchMode=yes",
        "-o",
        f"ConnectTimeout={CONNECT_TIMEOUT}",
        "-o",
        "StrictHostKeyChecking=accept-new",
        "-o",
        "ControlMaster=auto",
        "-o",
        f"ControlPath={control}/%C",
        "-o",
        "ControlPersist=600",
    ]


def agent_socket() -> str | None:
    """SSH_AUTH_SOCK, or a standard agent socket when it was not inherited, so
    hardware keys, 1Password, and ssh-agent work out of the box."""
    if "SSH_AUTH_SOCK" in os.environ:
        return os.environ["SSH_AUTH_SOCK"]
    for candidate in (
        os.path.expanduser("~/.ssh/agent.sock"),
        os.path.expanduser(
            "~/Library/Group Containers/2BUA8C4S2C.com.1password/t/agent.sock"
        ),
    ):
        if os.path.exists(candidate):
            return candidate
    return None


def env() -> dict[str, str]:
    """The local environment for ssh client processes."""
    environment = dict(os.environ)
    socket = agent_socket()
    if socket is not None:
        environment["SSH_AUTH_SOCK"] = socket
    return environment


def in_login_shell(script: str) -> str:
    """The remote command that runs `script` under the user's login shell.

    Non-interactive SSH commands run under `$SHELL -c` without sourcing
    `/etc/profile` or user profiles, which leaves PATH minimal (missing
    Nix, Homebrew, or user tool paths); a login shell finds python3 and the
    user's tools. Profiles often print (a motd, a greeting), which would
    corrupt a framed stdout, so the login shell starts with stdout on stderr
    and the real stdout kept on fd 3, and `script` gets it back. The outer
    `/bin/sh` makes the redirections work whatever the user's shell is.
    """
    login = f"exec 1>&3 3>&-; {script}"
    outer = f'exec 3>&1 1>&2; exec "${{SHELL:-/bin/sh}}" -l -c {shlex.quote(login)}'
    return f"/bin/sh -c {shlex.quote(outer)}"


def commands(target: str, home: str) -> dict[str, object]:
    """Everything the daemon runs on a host whose home is known, as complete
    remote commands: it never quotes anything itself. Variable inputs travel
    on stdin (the bridge's first frame, the ladder's request line, the run
    directory to remove), never in the command line."""
    bundle = f"{home}/{staged_name()}"
    python = "exec python3 -u "
    answer: dict[str, object] = {
        "host": target,
        "argv": [*base(), target],
        "control": f"{control_dir()}/%C",
        "bundle": bundle,
        "bridge": in_login_shell(
            "ALBEDO_JOB_ADMISSION=1 "
            + python
            + shlex.quote(f"{bundle}/albedo_bridge.py")
            + " --frame"
        ),
        "signal": in_login_shell(
            python + shlex.quote(f"{bundle}/albedo_signal.py") + " -"
        ),
        "remove": in_login_shell('read -r run && rm -rf -- "$run"'),
        "gather": in_login_shell(python + shlex.quote(f"{bundle}/albedo_gather.py")),
    }
    socket = agent_socket()
    if socket is not None:
        answer["auth_sock"] = socket
    return answer


def staged_name(digest: str | None = None) -> str:
    """The bundle's directory under the remote home."""
    return f"{STAGED_DIR}/{(digest or albedo_bundle.digest())[:16]}"


def archive() -> bytes:
    data = io.BytesIO()
    with tarfile.open(fileobj=data, mode="w") as tar:
        tar.add(
            str(albedo_bundle.ROOT),
            arcname=".",
            filter=lambda info: None if "__pycache__" in info.name else info,
        )
    return data.getvalue()


def stage_script(remote: str) -> str:
    """Unpack beside the target, then move into place, so a bundle directory
    that exists is always complete."""
    return (
        f'mkdir -p "{remote}.part.$$" && tar -C "{remote}.part.$$" -xf - && '
        f'{{ mv "{remote}.part.$$" "{remote}" 2>/dev/null || rm -rf "{remote}.part.$$"; }}'
    )


def run(
    target: str, script: str, timeout: float, data: bytes | None = None
) -> subprocess.CompletedProcess[bytes]:
    return subprocess.run(
        [*base(), target, script],
        input=data,
        stdin=None if data is not None else subprocess.DEVNULL,
        capture_output=True,
        timeout=timeout,
        env=env(),
    )


def failure(stderr: str) -> tuple[str, str]:
    """What ssh's own exit (255) means: a host we could reach with a person's
    help, or one we could not reach at all."""
    text = stderr.strip()
    lines = [line for line in text.splitlines() if line.strip()]
    detail = lines[-1] if lines else "ssh failed"
    lowered = text.lower()
    interactive = (
        "permission denied",
        "host key verification failed",
        "passphrase",
        "keyboard-interactive",
        "verification code",
        "too many authentication failures",
    )
    if any(phrase in lowered for phrase in interactive):
        return "needs_auth", detail
    return "unreachable", detail


PROBE = (
    "uname -s; uname -m; printf '%s\\n' \"$HOME\"; "
    "getconf _NPROCESSORS_ONLN 2>/dev/null || nproc 2>/dev/null || echo 1; "
    "python3 -c 'import sys; print(\"%d.%d\" % sys.version_info[:2])' 2>/dev/null "
    "|| echo none; "
    'test -f "$HOME/{bundle}/albedo_bridge.py" && echo staged || echo missing'
)


def probe(
    target: str, step: Callable[[str], None] = lambda _: None
) -> dict[str, object]:
    answer: dict[str, object] = {
        "host": target,
        "argv": [*base(), target],
        "control": f"{control_dir()}/%C",
    }
    bundle = staged_name()
    try:
        done = run(target, in_login_shell(PROBE.format(bundle=bundle)), PROBE_TIMEOUT)
    except subprocess.TimeoutExpired:
        return {**answer, "state": "unreachable", "detail": "ssh timed out"}
    except OSError as error:
        return {**answer, "state": "unreachable", "detail": f"ssh: {error}"}
    fields = done.stdout.decode(errors="replace").splitlines()
    if done.returncode == 255 or len(fields) != 6:
        state, detail = failure(done.stderr.decode(errors="replace"))
        return {**answer, "state": state, "detail": detail}
    system, arch, home, cpus, python, staged = (field.strip() for field in fields)
    answer.update(
        commands(target, home),
        os=system,
        arch=arch,
        home=home,
        cpus=int(cpus) if cpus.isdigit() and int(cpus) > 0 else 1,
        python=python,
    )
    version = (
        tuple(int(part) for part in python.split(".")) if python[:1].isdigit() else None
    )
    if version is None:
        return {
            **answer,
            "state": "unsupported",
            "detail": "needs python >= 3.11 (python3 not found)",
        }
    if version < MINIMUM_PYTHON:
        return {
            **answer,
            "state": "unsupported",
            "detail": f"needs python >= 3.11 (found {python})",
        }
    if staged != "staged":
        step("staging")
        try:
            staging = run(
                target,
                in_login_shell(stage_script(f"$HOME/{bundle}")),
                STAGE_TIMEOUT,
                archive(),
            )
        except subprocess.TimeoutExpired:
            return {**answer, "state": "unreachable", "detail": "staging timed out"}
        if staging.returncode != 0:
            detail = staging.stderr.decode(errors="replace").strip()[-500:]
            return {
                **answer,
                "state": "unreachable",
                "detail": f"staging failed: {detail}",
            }
    return {**answer, "state": "ready", "detail": ""}


def config_hosts(path: str | None = None, seen: set[str] | None = None) -> list[str]:
    """`Host` names in ~/.ssh/config and the files it includes, in order,
    without patterns (`*`, `?`, `!`): the hosts a person can name."""
    path = path or os.path.expanduser("~/.ssh/config")
    seen = set() if seen is None else seen
    if path in seen:
        return []
    seen.add(path)
    try:
        with open(path, encoding="utf-8", errors="replace") as handle:
            lines = handle.read().splitlines()
    except OSError:
        return []
    names: list[str] = []
    for line in lines:
        words = shlex.split(line, comments=True) if line.strip() else []
        if len(words) < 2:
            continue
        keyword = words[0].lower()
        if keyword == "host":
            names += [w for w in words[1:] if not any(c in w for c in "*?!")]
        elif keyword == "include":
            for pattern in words[1:]:
                pattern = os.path.expanduser(pattern)
                if not os.path.isabs(pattern):
                    pattern = os.path.join(os.path.expanduser("~/.ssh"), pattern)
                for included in sorted(glob.glob(pattern)):
                    names += config_hosts(included, seen)
    return list(dict.fromkeys(names))


def announce(step: str) -> None:
    print(json.dumps({"step": step}), flush=True)


def main(argv: list[str]) -> int:
    if len(argv) == 3 and argv[1] == "probe":
        print(json.dumps(probe(argv[2], announce)))
        return 0
    if len(argv) == 2 and argv[1] == "hosts":
        print(json.dumps(config_hosts()))
        return 0
    if len(argv) == 4 and argv[1] == "commands":
        print(json.dumps(commands(argv[2], argv[3])))
        return 0
    print(
        "usage: albedo_ssh.py probe <[user@]host> | commands <[user@]host> <home>",
        file=sys.stderr,
    )
    return 2


if __name__ == "__main__":
    sys.exit(main(sys.argv))