import modal

app = modal.App("surflo-ctrl")
VOL = modal.Volume.from_name("surflo-vol", create_if_missing=True)

img = (
    modal.Image.from_registry("nvidia/cuda:12.4.1-devel-ubuntu22.04", add_python="3.10")
    .apt_install("software-properties-common", "curl", "git", "build-essential",
                 "cmake", "ninja-build")
    .run_commands(
        "add-apt-repository -y ppa:deadsnakes/ppa",
        "apt-get update -y",
        "apt-get install -y python3.10 python3.10-venv python3.10-dev python3.10-distutils",
    )
)

PAYLOAD = r'''
import os, subprocess, sys, time

VENV = "/vol/venv"
PY = VENV + "/bin/python"
REPO = "/vol/Surflo"
CKPT = "/vol/checkpoints/surflo_v0.pt"
SAMPLE = REPO + "/media/sample"
log = []


def sh(cmd, label, timeout=7200):
    t0 = time.time()
    env = dict(os.environ)
    env.pop("PYTHONPATH", None)
    try:
        p = subprocess.run(["bash", "-lc", cmd], capture_output=True, text=True,
                           timeout=timeout, env=env, cwd="/")
    except subprocess.TimeoutExpired:
        log.append("=== %s TIMEOUT" % label)
        return 1
    dt = time.time() - t0
    log.append("=== %s rc=%d %.0fs\n%s" % (label, p.returncode, dt,
               ((p.stdout or "") + (p.stderr or ""))[-3000:]))
    print("=== %s rc=%d %.0fs" % (label, p.returncode, dt), flush=True)
    return p.returncode


def save():
    open("/vol/ctrl_run.txt", "w").write("\n".join(log))


# --- wait for the build to have produced a working stack -------------------
READY = """
import importlib, sys
need = ['torch','diff_gaussian_rasterization_surflo','fused_ssim','geodel','nvdiffrast.torch']
bad = []
for m in need:
    try:
        importlib.import_module(m)
    except Exception as e:
        bad.append(m)
print('READY' if not bad else 'NOT_READY ' + ','.join(bad))
"""
for attempt in range(60):
    if os.path.exists("/vol/build310.txt"):
        log.append("build310.txt present")
        break
    log.append("wait_build %d" % attempt)
    print("wait_build %d" % attempt, flush=True)
    time.sleep(60)

for attempt in range(20):
    r = subprocess.run([PY, "-c", READY], capture_output=True, text=True)
    out = (r.stdout or "") + (r.stderr or "")
    log.append("wait %d: %s" % (attempt, out.strip()[:200]))
    print("wait %d: %s" % (attempt, out.strip()[:120]), flush=True)
    if "READY" in out and "NOT_READY" not in out:
        break
    time.sleep(60)
else:
    log.append("!! stack never became ready")
    save()
    sys.exit("not ready")

# --- VALID CONTROL: guided on Surflo's own real photos ---------------------
n = len([f for f in os.listdir(SAMPLE) if f.lower().endswith((".jpg", ".jpeg", ".png"))])
log.append("sample images: %d" % n)

OUT = "/vol/outputs/ctrl_guided"
sh("cd %s && %s scripts/infer.py mode=guided guided=default ckpt=%s "
   "source.image_folder=%s source.n_images=16 num_query_points=100000 output_dir=%s"
   % (REPO, PY, CKPT, SAMPLE, OUT), "ctrl_guided_default", timeout=7200)

listing = ""
for root, _d, files in os.walk(OUT):
    for f in sorted(files):
        listing += "    %s %s\n" % (os.path.join(root, f),
                                    "{:,}".format(os.path.getsize(os.path.join(root, f))))
log.append("--- outputs ---\n" + listing)
save()
'''

import subprocess


@app.function(timeout=28800, volumes={"/vol": VOL}, image=img, gpu="L4")
def run_ctrl() -> str:
    import subprocess as sp
    with open("/vol/_ctrl.py", "w") as fh:
        fh.write(PAYLOAD)
    p = sp.run(["bash", "-lc", "/usr/bin/python3.10 /vol/_ctrl.py"],
               capture_output=True, text=True, timeout=28800)
    VOL.commit()
    return ((p.stdout or "") + (p.stderr or ""))[-4000:]


@app.local_entrypoint()
def main():
    print(run_ctrl.remote())