#!/usr/bin/env python3
"""Birth one clip per section, from one portrait plus approved audio, on a rented GPU.

    python3 nacer.py --portrait retrato-verde.png --audio parte-*.wav [--scale 2.5] [--match -8.5]

One clip per audio file, no fifteen-second cap, no seams, the same face throughout.
Outputs land in ./salidas/, loudness-matched, with the mouth measured against the house band.
The machine is always returned: it terminates on success, on failure and on Ctrl-C.
Cost: about $33-37 for one published minute of video (A100 80GB at $1.39-1.59/hour,
~2.5 h per 6.4 s clip un-tuned). Corrected 31-08-2026 — was wrongly given as $2.5/hour.
"""
import argparse, json, os, re, signal, subprocess, sys, time, urllib.request, urllib.error

BASE = os.path.dirname(os.path.abspath(__file__))
KEYS = os.path.expanduser("~/Claude/6 - Private/keys/keys.env")
API = "https://rest.runpod.io/v1"
IMAGE = "runpod/pytorch:1.1.0-cu1281-torch260-ubuntu2204"
# Cards, cheapest first. MEASURED 01-09-2026 (audit round, pod kfa3q47kh4o8rf): the 14B model
# LOADS on a 48 GB card, unquantised, at $0.44 to $0.478 an hour, which puts a published minute
# near $10 against $33-37 on an A100. The August "dies on container memory" was InfiniteTalk,
# a different engine, not this one. Which of the five cards actually came is NOT recorded, so
# the rule is "a 48 GB card from this list", never a named one. --gpu overrides.
GPUS = ["NVIDIA A40", "NVIDIA L40", "NVIDIA L40S", "NVIDIA RTX A6000",
        "NVIDIA RTX 6000 Ada Generation", "NVIDIA A100 80GB PCIe", "NVIDIA A100-SXM4-80GB"]
PROMPT = ("A person speaks warmly and calmly straight to the camera, subtle natural head "
          "movement, flat solid green background, even studio lighting")

def log(m):
    print(time.strftime("%H:%M:%S ") + m, flush=True)

def key():
    for l in open(KEYS):
        m = re.match(r'^(?:export )?RUNPOD_API_KEY=(.*)$', l.strip())
        if m: return m.group(1).strip().strip('"').strip("'")
    sys.exit("no RUNPOD_API_KEY in " + KEYS)

APIKEY = key()
POD = {"id": None, "rented": False}

def rp(method, path, body=None):
    req = urllib.request.Request(API + path, method=method,
        headers={"Authorization": f"Bearer {APIKEY}", "Content-Type": "application/json"},
        data=json.dumps(body).encode() if body is not None else None)
    try:
        return json.loads(urllib.request.urlopen(req, timeout=60).read().decode() or "{}")
    except urllib.error.HTTPError as e:
        log(f"HTTP {e.code} {method} {path}: {e.read().decode()[:300]}"); raise

def terminate(*_):
    if POD["id"]:
        try:
            rp("DELETE", f"/pods/{POD['id']}"); log(f"pod {POD['id']} terminated")
        except Exception as e:
            log(f"COULD NOT TERMINATE {POD['id']}: {e}")
        POD["id"] = None

def sh(cmd, timeout=180):
    return subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout)

def balance():
    """What is left on the account, so a run cannot die halfway for lack of credit."""
    req = urllib.request.Request("https://api.runpod.io/graphql", method="POST",
        headers={"Authorization": f"Bearer {APIKEY}", "Content-Type": "application/json"},
        data=json.dumps({"query": "query { myself { clientBalance } }"}).encode())
    try:
        return json.loads(urllib.request.urlopen(req, timeout=30).read())["data"]["myself"]["clientBalance"]
    except Exception:
        return None

def ours():
    """Every machine this tool family may have left behind."""
    try:
        return [p for p in rp("GET", "/pods") or []
                if str(p.get("name", "")).startswith(("nacer-", "valentina-", "speaker-"))]
    except Exception as e:
        log(f"could not list pods: {e}"); return []

def sweep(force=False):
    """Kill every machine of ours. Run it before and after any GPU work: an orphan
    bills quietly for hours. One did, on 31-08-2026, for $2.25."""
    left = ours()
    if not left:
        log("sweep: no machine of ours is renting"); return 0
    for p in left:
        log(f"sweep: terminating {p.get('id')} ({p.get('name')}, {p.get('desiredStatus')})")
        try: rp("DELETE", f"/pods/{p['id']}")
        except Exception as e: log(f"  COULD NOT TERMINATE: {e}")
    return len(left)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--portrait", help="green-screen portrait, PNG")
    ap.add_argument("--audio", nargs="+", help="one WAV per section (16 kHz mono is made for you)")
    ap.add_argument("--scale", default="2.5", help="how strongly the mouth follows the sound (2.5 soft, 3.0 stronger)")
    ap.add_argument("--steps", default="40")
    ap.add_argument("--match", default="-8.5", help="LUFS to match; the engine returns the voice ~14 dB down")
    ap.add_argument("--out", default=os.path.join(os.getcwd(), "salidas"))
    ap.add_argument("--max-usd", type=float, default=12.0,
                    help="hard ceiling: the machine is returned when the run has cost this much")
    ap.add_argument("--rate", type=float, default=1.39, help="$/hour of the machine, for the ceiling")
    ap.add_argument("--sweep", action="store_true", help="terminate every machine of ours and exit")
    ap.add_argument("--disk", type=int, default=140,
                    help="container disk in GB, and it is billed: the weights need about 110")
    ap.add_argument("--gpu", nargs="+", help="cards to ask for, cheapest first; overrides the default list")
    a = ap.parse_args()

    if a.sweep:
        sweep(); return
    if not (a.portrait and a.audio):
        sys.exit("give --portrait and --audio (or --sweep on its own to return every machine of ours)")

    # Nothing is rented before these three answers.
    left = ours()
    if left:
        sys.exit(f"a machine of ours is already renting ({', '.join(p['id'] for p in left)}). "
                 f"Use it, or run: python3 nacer.py --sweep --portrait x --audio y")
    bal = balance()
    if bal is not None:
        log(f"balance ${bal:.2f}, ceiling for this run ${a.max_usd:.2f}")
        if bal < a.max_usd:
            sys.exit(f"balance ${bal:.2f} is under the ceiling ${a.max_usd:.2f}: top up first, "
                     f"or lower --max-usd. A run that dies halfway pays for nothing.")

    os.makedirs(a.out, exist_ok=True)
    prep = os.path.join(a.out, "entradas"); os.makedirs(prep, exist_ok=True)
    parts = []
    for i, w in enumerate(a.audio, 1):
        pid = f"parte{i:02d}"
        dst = os.path.join(prep, pid + ".wav")
        subprocess.run(["ffmpeg","-y","-v","error","-i",w,"-ac","1","-ar","16000","-c:a","pcm_s16le",dst], check=True)
        parts.append({"id": pid, "audio": pid + ".wav", "source": w})
    jobs = {"prompt": PROMPT, "portrait": os.path.basename(a.portrait),
            "scale": float(a.scale), "steps": int(a.steps), "parts": parts}
    json.dump(jobs, open(os.path.join(a.out, "jobs.json"), "w"), indent=1)
    log(f"{len(parts)} section(s) to birth at scale {a.scale}")

    signal.signal(signal.SIGINT, lambda *x: (terminate(), sys.exit(1)))
    pub = open(os.path.expanduser("~/.ssh/id_ed25519.pub")).read().strip()
    pod = rp("POST", "/pods", {"name": "nacer-secciones", "imageName": IMAGE, "gpuTypeIds": (a.gpu or GPUS),
        "cloudType": "SECURE", "containerDiskInGb": a.disk, "ports": ["22/tcp"], "env": {"PUBLIC_KEY": pub}})
    POD["id"] = pod.get("id"); POD["rented"] = True; log(f"pod {POD['id']}")
    # The rate comes from the machine we were GIVEN, never from a number carried in. Asking for
    # several cards means RunPod hands over whatever is free: an inherited $0.44 (true of an A40)
    # was applied to an L40S that bills $0.99, and every forecast built on it was half the real
    # price (Valentina seat, 01-09-2026, the third inherited-number fault in two days).
    real = pod.get("costPerHr")
    if real:
        real = float(real)
        log(f"the machine given bills ${real:.2f}/h (asked: {', '.join(a.gpu or GPUS)})")
        if abs(real - a.rate) > 0.05:
            log(f"NOTE: ceiling was set at ${a.rate:.2f}/h; using the machine's ${real:.2f}/h from here")
            a.rate = real
    else:
        log(f"the machine did not report its price; the ceiling uses ${a.rate:.2f}/h, a GUESS")

    ip = port = None; t0 = time.time()
    while time.time() - t0 < 1800:
        time.sleep(20)
        info = rp("GET", f"/pods/{POD['id']}")
        ip = info.get("publicIp"); port = (info.get("portMappings") or {}).get("22")
        if ip and port: break
    if not (ip and port):
        terminate(); sys.exit("no ssh mapping after 30 min")

    SSH = f"ssh -o StrictHostKeyChecking=no -o ConnectTimeout=15 -p {port} root@{ip}"
    SCP = f"scp -o StrictHostKeyChecking=no -P {port}"
    for _ in range(40):
        if sh(f"{SSH} 'echo ok'").stdout.strip() == "ok": break
        time.sleep(15)
    else:
        terminate(); sys.exit("ssh never came up")
    log("ssh up, sending portrait and audio")

    sh(f"{SSH} 'mkdir -p /root/inputs /root/out'")
    sh(f"{SCP} '{a.portrait}' root@{ip}:/root/inputs/", timeout=300)
    sh(f"{SCP} {' '.join(repr(os.path.join(prep, p['audio'])) for p in parts)} root@{ip}:/root/inputs/", timeout=600)
    sh(f"{SCP} '{os.path.join(a.out,'jobs.json')}' root@{ip}:/root/jobs.json", timeout=120)
    sh(f"{SCP} '{BASE}/remote-base.sh' root@{ip}:/root/remote.sh", timeout=120)
    try:
        sh(f"{SSH} 'chmod +x /root/remote.sh; nohup bash /root/remote.sh >/dev/null 2>&1 </dev/null & disown; echo go'", timeout=30)
    except subprocess.TimeoutExpired:
        pass
    if not sh(f"{SSH} 'sleep 3; cat /root/status.txt 2>/dev/null'").stdout.strip():
        terminate(); sys.exit("remote script never started")
    log("running on the machine")

    ok = False
    t0 = time.time()
    while time.time() - t0 < 12 * 3600:
        time.sleep(300)
        spent = (time.time() - t0) / 3600 * a.rate
        if spent > a.max_usd:
            log(f"ceiling reached (${spent:.2f} > ${a.max_usd:.2f}): returning the machine")
            break
        try:
            st = sh(f"{SSH} 'cat /root/status.txt; ls /root/DONE /root/FAILED 2>/dev/null'").stdout.replace("\n", " ").strip()
        except Exception as e:
            # a hung poll must never escape: it would skip the terminate below and orphan the machine
            log(f"poll failed ({type(e).__name__}), the machine stays owned by this script"); continue
        log(f"status: {st}  (${spent:.2f} spent)")
        if "FAILED" in st: break
        if "DONE" in st: ok = True; break

    sh(f"{SCP} -r root@{ip}:/root/out/. '{a.out}/'", timeout=1800)
    sh(f"{SCP} root@{ip}:/root/build.log '{a.out}/build.log'", timeout=300)
    sh(f"{SCP} root@{ip}:/root/card.txt '{a.out}/card.txt'", timeout=120)
    try: log("card served: " + open(f"{a.out}/card.txt").read().strip().replace("\n", " | "))
    except Exception: log("card not recorded")
    terminate()
    if not ok:
        sys.exit(f"the machine returned nothing; read {a.out}/build.log")

    # the engine returns the voice about 14 dB down: put it back where the page expects it
    for f in sorted(os.listdir(a.out)):
        if not f.endswith(".mp4"): continue
        src = os.path.join(a.out, f); tmp = src + ".lvl.mp4"
        subprocess.run(["ffmpeg","-y","-v","error","-i",src,"-c:v","copy",
                        "-af",f"loudnorm=I={a.match}:TP=-1.0:LRA=11","-c:a","aac",tmp], check=True)
        os.replace(tmp, src)
        log(f"{f} levelled to {a.match} LUFS")
    log(f"done. Clips in {a.out}. Measure the mouth with medir-boca.py, then key and wire with /site-speaker.")

if __name__ == "__main__":
    try:
        main()
    finally:
        terminate()
        # belt and braces: if anything of ours survived, it dies here
        try: sweep()
        except Exception as e: log(f"final sweep failed: {e}")