#!/usr/bin/env python3
"""
One-shot deploy script for the Heicode Manager Azure VM. Designed to be
idempotent — running it twice in a row converges to the same state.

What it does:
  1. SSH into the VM with credentials supplied via env (NEVER hard-coded
     in the repo — the launcher passes them on stdin).
  2. cd into the manager checkout, git pull origin main.
  3. docker compose pull + up -d --build, then prune dangling images.
  4. Wait for /api/status to return 200.
  5. Print container status + tail of the manager logs.

Failure mode: if any step exits non-zero, the whole script aborts and
prints the offending command's stderr — caller (this Claude session)
escalates to the user instead of silently moving on.
"""

import os
import sys
import time
import paramiko

HOST = os.environ.get("MGR_HOST", "20.24.50.121")
USER = os.environ.get("MGR_USER", "heicode")
PASS = os.environ.get("MGR_PASS", "")
PORT = int(os.environ.get("MGR_PORT", "22"))

if not PASS:
    print("MGR_PASS not set", file=sys.stderr)
    sys.exit(2)


def run(client: paramiko.SSHClient, cmd: str, timeout: int = 600) -> int:
    """Stream a remote command's output to local stdout/stderr."""
    print(f"\n$ {cmd}", flush=True)
    stdin, stdout, stderr = client.exec_command(cmd, timeout=timeout, get_pty=True)
    for line in iter(stdout.readline, ""):
        if not line:
            break
        sys.stdout.write(line)
        sys.stdout.flush()
    rc = stdout.channel.recv_exit_status()
    err = stderr.read().decode("utf-8", errors="replace")
    if err.strip():
        sys.stderr.write(err)
    return rc


def must(rc: int, label: str) -> None:
    if rc != 0:
        print(f"\nFAIL: {label} exit={rc}", file=sys.stderr)
        sys.exit(rc or 1)


def main() -> None:
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    print(f"Connecting {USER}@{HOST}:{PORT} ...", flush=True)
    client.connect(
        hostname=HOST,
        port=PORT,
        username=USER,
        password=PASS,
        look_for_keys=False,
        allow_agent=False,
        timeout=30,
    )

    # Locate the manager checkout. We try a few likely paths instead of
    # hardcoding one — first match wins.
    probe = (
        "for p in "
        "~/heicode-mananger ~/heicode ~/heicode-mananger/heicode "
        "/opt/heicode-mananger /home/heicode/heicode-mananger; do "
        "  if [ -d $p/.git ]; then echo FOUND=$p; break; fi; "
        "done"
    )
    stdin, stdout, _ = client.exec_command(probe, timeout=20)
    out = stdout.read().decode().strip()
    print(out)
    repo = None
    for line in out.splitlines():
        if line.startswith("FOUND="):
            repo = line.split("=", 1)[1]
            break
    if not repo:
        print("Could not locate manager git checkout on VM", file=sys.stderr)
        sys.exit(3)
    print(f"Using repo: {repo}")

    # Find docker-compose file we'll target.
    compose = "docker-compose.azure-vm.yml"
    must(run(client, f"test -f {repo}/heicode/{compose} || ls {repo}/heicode/docker-compose*.yml"),
         "locate compose file")

    # Pull latest code.
    must(run(client, f"cd {repo} && git fetch origin && git checkout main && git pull --ff-only origin main"),
         "git pull")

    # Show what we have locally now.
    run(client, f"cd {repo} && git log --oneline -3")

    # Rebuild + restart.
    must(run(client,
             f"cd {repo}/heicode && sudo docker compose -f {compose} up -d --build 2>&1",
             timeout=1800),
         "compose up")

    # Clean dangling images per CLAUDE.md.
    run(client, "sudo docker image prune -f")

    # Wait for /api/status to come back.
    print("\nWaiting for /api/status to return 200 (max 90s) ...")
    deadline = time.time() + 90
    status_ok = False
    while time.time() < deadline:
        rc = run(client,
                 "curl -fsS -o /dev/null -w 'HTTP_%{http_code}\\n' https://code.xinghanlab.com/api/status")
        if rc == 0:
            status_ok = True
            break
        time.sleep(5)
    if not status_ok:
        print("FAIL: /api/status did not return 200 within 90s", file=sys.stderr)
        run(client, f"cd {repo}/heicode && sudo docker compose -f {compose} logs --tail=80 heicode")
        sys.exit(4)

    print("\n--- Final state ---")
    run(client, "sudo docker ps --format 'table {{.Names}}\\t{{.Status}}\\t{{.Image}}'")
    run(client, f"cd {repo}/heicode && sudo docker compose -f {compose} logs --tail=20 heicode")

    client.close()
    print("\nDEPLOY OK")


if __name__ == "__main__":
    main()