#!/usr/bin/env python3
"""
SSH to a Linux host (where Docker + redroid run) and run diagnostics or shell commands.

Authentication (pick one): private key (--key-file / REDROID_SSH_KEY), password
(REDROID_SSH_PASSWORD), or --password-prompt. Do not commit secrets.

Examples:
  REDROID_SSH_HOST=192.168.2.179 REDROID_SSH_USER=eric python tools/ssh_remote.py --password-prompt diagnose

  python tools/ssh_remote.py -H 192.168.2.179 -U eric --key-file ~/.ssh/id_ed25519 diagnose

  python tools/ssh_remote.py -H ... -U eric --password-prompt exec -- bash -lc 'docker ps -a | head'
"""

from __future__ import annotations

import argparse
import getpass
import os
import shlex
import sys
import threading
from pathlib import Path

_REPO_ROOT = Path(__file__).resolve().parents[1]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))

import paramiko


def connect_client(args: argparse.Namespace) -> paramiko.SSHClient:
    host = (args.host or "").strip()
    user = (args.user or "").strip()
    if not host or not user:
        print("Set -H/--host and -U/--user (or REDROID_SSH_HOST / REDROID_SSH_USER).", file=sys.stderr)
        sys.exit(2)

    password = (getattr(args, "password", None) or os.environ.get("REDROID_SSH_PASSWORD", "")).strip() or None
    key_file = (args.key_file or os.environ.get("REDROID_SSH_KEY", "")).strip() or None

    if getattr(args, "password_prompt", False):
        password = getpass.getpass("SSH password: ") or None
        if not password:
            print("Empty password.", file=sys.stderr)
            sys.exit(2)

    if not password and not key_file:
        print(
            "Need auth: --key-file / REDROID_SSH_KEY, or REDROID_SSH_PASSWORD, or --password-prompt.",
            file=sys.stderr,
        )
        sys.exit(2)

    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    kw: dict = {
        "hostname": host,
        "port": int(args.port),
        "username": user,
        "timeout": float(args.timeout),
        "allow_agent": bool(key_file),
        "look_for_keys": bool(key_file),
    }
    if key_file:
        kw["key_filename"] = os.path.expanduser(key_file)
    if password:
        kw["password"] = password

    client.connect(**kw)
    return client


def run_remote(client: paramiko.SSHClient, bash_script: str, timeout: float = 300.0) -> int:
    cmd = "bash -lc {}".format(shlex.quote(bash_script))
    stdin, stdout, stderr = client.exec_command(cmd)
    stdout.channel.settimeout(timeout)
    out = stdout.read().decode("utf-8", errors="replace")
    err = stderr.read().decode("utf-8", errors="replace")
    if out:
        sys.stdout.write(out)
        if not out.endswith("\n"):
            sys.stdout.write("\n")
    if err:
        sys.stderr.write(err)
        if not err.endswith("\n"):
            sys.stderr.write("\n")
    return int(stdout.channel.recv_exit_status())


def cmd_diagnose(args: argparse.Namespace) -> int:
    cname = args.container_name
    port = int(args.adb_port)
    qc = shlex.quote(cname)
    script = """
set +e
echo "=== docker ps (name filter) ==="
docker ps -a --filter name={qc} --no-trunc 2>/dev/null | head -35
echo "=== docker inspect ==="
docker inspect --format 'status={{{{.State.Status}}}} exit={{{{.State.ExitCode}}}} err={{{{.State.Error}}}}' {qc} 2>/dev/null || echo "(no such container)"
echo "=== docker logs (tail 150) ==="
docker logs --tail 150 {qc} 2>/dev/null || echo "(no logs)"
echo "=== adb ==="
adb connect 127.0.0.1:{port}
adb devices -l
adb -s 127.0.0.1:{port} shell getprop sys.boot_completed 2>/dev/null || true
echo "=== adb logcat (display/GLES, last 80) ==="
adb -s 127.0.0.1:{port} logcat -d -t 80 '*:S' SurfaceFlinger:V EGL:V OpenGLRenderer:V redroid:V AndroidRuntime:E 2>/dev/null || true
""".format(
        qc=qc,
        port=port,
    )
    client = connect_client(args)
    try:
        return run_remote(client, script)
    finally:
        client.close()


def cmd_exec(args: argparse.Namespace) -> int:
    remote = " ".join(args.remote_words).strip()
    if not remote:
        print("exec: provide a remote command after --", file=sys.stderr)
        return 2
    client = connect_client(args)
    try:
        return run_remote(client, remote, timeout=float(args.timeout))
    finally:
        client.close()


def cmd_shell(args: argparse.Namespace) -> int:
    """Minimal interactive shell (no full PTY); useful for quick typing. Ctrl+D to exit."""
    client = connect_client(args)
    try:
        chan = client.invoke_shell()
        chan.send("export TERM=dumb\n")

        def read_loop():
            try:
                while True:
                    data = chan.recv(4096)
                    if not data:
                        break
                    sys.stdout.buffer.write(data)
                    sys.stdout.buffer.flush()
            except EOFError:
                pass

        t = threading.Thread(target=read_loop, daemon=True)
        t.start()
        try:
            while True:
                line = sys.stdin.readline()
                if not line:
                    break
                chan.send(line)
        except KeyboardInterrupt:
            chan.close()
        return 0
    finally:
        client.close()


def main() -> None:
    ap = argparse.ArgumentParser(description="SSH helper for remote Docker/adb diagnostics.")
    ap.add_argument("-H", "--host", default=os.environ.get("REDROID_SSH_HOST", ""))
    ap.add_argument("-U", "--user", default=os.environ.get("REDROID_SSH_USER", ""))
    ap.add_argument("-P", "--port", type=int, default=int(os.environ.get("REDROID_SSH_PORT", "22")))
    ap.add_argument("--key-file", default=os.environ.get("REDROID_SSH_KEY", ""))
    ap.add_argument(
        "--password",
        default="",
        help="Unsafe on shell history; prefer REDROID_SSH_PASSWORD or --password-prompt.",
    )
    ap.add_argument("--password-prompt", action="store_true")
    ap.add_argument("--timeout", type=float, default=60.0, help="SSH connect / exec timeout (seconds)")

    sub = ap.add_subparsers(dest="command", required=True)

    d = sub.add_parser("diagnose", help="docker ps/inspect/logs + adb on the remote host")
    d.add_argument(
        "--container-name",
        default=os.environ.get("REDROID_CONTAINER", "redroid_13.0.0_mindthegapps_houdini_magisk"),
    )
    d.add_argument("--adb-port", type=int, default=int(os.environ.get("REDROID_ADB_PORT", "5570")))

    e = sub.add_parser("exec", help="Run one remote bash -lc string (pass command after --)")
    e.add_argument("remote_words", nargs=argparse.REMAINDER, help="e.g. -- bash -lc 'docker ps'")

    s = sub.add_parser("shell", help="Interactive SSH shell (basic; prefer OpenSSH ssh for full TTY)")

    ns = ap.parse_args()
    if ns.command == "diagnose":
        sys.exit(cmd_diagnose(ns))
    if ns.command == "exec":
        sys.exit(cmd_exec(ns))
    if ns.command == "shell":
        sys.exit(cmd_shell(ns))
    sys.exit(2)


if __name__ == "__main__":
    main()