#!/usr/bin/env python3
"""Railway free-VM farm manager.

Streams keyhunter-imba/railway/run_node.sh into `ssh railway.new` (one
persistent free VM per source IP). The VM clones a fleet node repo, runs
run_imba.py and pushes results/ back — fleet_collect.py harvests it as a
normal node run.

Usage:
  python railway_farm.py --node imba-node-0980
  python railway_farm.py --repo pro2628/imba-node-0980 --budget-min 480
  python railway_farm.py --node imba-node-0980 --dry
"""
import argparse
import json
import os
import signal
import subprocess
import sys
import tempfile
import time
from pathlib import Path

PROJ = Path(__file__).parent
DEPLOY = Path(os.environ.get("KH_DEPLOY", str(PROJ / "data" / "fleet_deploy.json")))
RUNNER = Path(os.environ.get("KH_RUNNER",
                             str(PROJ.parent / "keyhunter-imba" / "railway" / "run_node.sh")))
SSH = ["ssh", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=20",
       "-o", "ServerAliveInterval=30", "-o", "ServerAliveCountMax=6",
       "railway.new", "bash -s"]


PREFERRED_LOGINS = {"pro2065", "pro2628"}


def find_rec(args):
    recs = json.load(open(DEPLOY, encoding="utf-8"))
    if args.node:
        matches = [r for r in recs if r.get("repo") == args.node and r.get("ok")]
        if not matches:
            sys.exit(f"node not found/ok: {args.node}")
        preferred = [r for r in matches if r.get("login") in PREFERRED_LOGINS]
        return preferred[0] if preferred else matches[0]
    if args.repo:
        owner, name = args.repo.split("/", 1)
        for r in recs:
            if r.get("login") == owner and r.get("repo") == name:
                return r
        sys.exit(f"repo not in fleet_deploy.json: {args.repo}")
    sys.exit("need --node NAME or --repo owner/name")


def build_payload(rec, budget, cycle_sleep):
    body = RUNNER.read_text(encoding="utf-8")
    body = body.replace("\r\n", "\n")  # Windows host must not ship CRLF
    head = (f"KH_PAT='{rec['pat']}'\n"
            f"KH_REPO='{rec['login']}/{rec['repo']}'\n"
            f"KH_BUDGET_MIN={budget}\n"
            f"KH_CYCLE_SLEEP={cycle_sleep}\n"
            "export KH_PAT KH_REPO KH_BUDGET_MIN KH_CYCLE_SLEEP\n\n")
    return head + body


def redact(line):
    if "KH_PAT=" in line:
        i = line.index("KH_PAT=")
        line = line[:i] + "KH_PAT=<redacted>" + line[i:].split("'", 2)[-1]
    if "ghp_" in line:
        i = line.index("ghp_")
        line = line[:i] + "ghp_<redacted>" + line[i:].split("ghp_")[-1][36:]
    return line


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--node", help="repo field from fleet_deploy.json")
    ap.add_argument("--repo", help="owner/name")
    ap.add_argument("--budget-min", type=int, default=240)
    ap.add_argument("--cycle-sleep", type=int, default=300)
    ap.add_argument("--log", default=str(PROJ / "data" / "railway_farm.log"))
    ap.add_argument("--dry", action="store_true")
    args = ap.parse_args()

    if not RUNNER.exists():
        sys.exit(f"runner missing: {RUNNER}")
    rec = find_rec(args)
    payload = build_payload(rec, args.budget_min, args.cycle_sleep)

    if args.dry:
        for line in payload.splitlines()[:3]:
            print(redact(line))
        print(f"payload: {len(payload)} bytes, dry run — no ssh")
        return

    logf = open(args.log, "a", encoding="utf-8", buffering=1)

    def log(msg):
        stamp = time.strftime("%Y-%m-%d %H:%M:%S")
        line = f"[{stamp}] {redact(msg)}"
        print(line, flush=True)
        logf.write(line + "\n")

    log(f"farm start node={rec['login']}/{rec['repo']} "
        f"budget={args.budget_min}m")

    proc = None
    attempt = 0
    running = True

    def stop(*_):
        nonlocal running
        running = False
        if proc and proc.poll() is None:
            proc.terminate()

    signal.signal(signal.SIGINT, stop)
    signal.signal(signal.SIGTERM, stop)

    while running:
        with tempfile.NamedTemporaryFile("w", suffix=".sh", delete=False,
                                         newline="\n") as tf:
            tf.write(payload)
            payload_path = tf.name
        attempt += 1
        log(f"ssh attempt {attempt}")
        try:
            with open(payload_path, "r") as stdin:
                proc = subprocess.Popen(
                    SSH, stdin=stdin, stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT, text=True,
                    encoding="utf-8", errors="replace", bufsize=1)
            assert proc.stdout is not None
            for line in proc.stdout:
                log(line.rstrip())
            proc.wait()
            log(f"ssh exited rc={proc.returncode}")
        except Exception as e:
            log(f"ssh error: {e}")
        finally:
            try:
                import os
                os.unlink(payload_path)
            except OSError:
                pass
        if not running:
            break
        backoff = min(60 * attempt, 600)
        log(f"RECONNECT in {backoff}s")
        time.sleep(backoff)

    log("farm stop")


if __name__ == "__main__":
    main()