#!/usr/bin/env python3
"""
lingbot_gen.py — run LingBot-World generation (image/footage + prompt →
explorable video world) against an installed repo + downloaded checkpoint.
Runs OUTSIDE the lingbot venv (any ambient python, stdlib-only) and invokes
the venv's python on the repo's own generate.py / generate_fast.py with the
arguments that script actually defines (--task/--size/--ckpt_dir/--image/
--prompt/--frame_num/... — read from the repo, not guessed).

Honest preflights, each with its own machine-readable failure code:
  - flash_attn is a stub → generation cannot run (the DiT hard-requires the
    real kernels); fails fast with the reason instead of a deep traceback.
  - checkpoint dir missing → tells the caller to `zelpi lingbot pull`.
  - low VRAM → auto-adds the repo's own low-memory knobs
    (--offload_model True --t5_cpu --convert_model_dtype) and warns; these
    models are 27B-class MoE, a single small GPU is likely still not enough.

Footage input: --footage <video> extracts the first frame (via the venv's
own opencv, which requirements.txt guarantees) as the conditioning image —
LingBot-World conditions on a single frame + prompt + optional camera path.
"""
from __future__ import annotations

import argparse
import json
import pathlib
import subprocess
import sys


def _venv_python(install_path: pathlib.Path) -> pathlib.Path:
    if sys.platform == "win32":
        return install_path / ".venv" / "Scripts" / "python.exe"
    return install_path / ".venv" / "bin" / "python"


def _fail(code: str, message: str):
    print("LINGBOT_GEN_RESULT " + json.dumps({"code": code, "message": message}))
    sys.exit(1)


def _extract_frame(venv_py: pathlib.Path, footage: pathlib.Path, out_png: pathlib.Path) -> bool:
    code = (
        "import cv2,sys\n"
        f"cap=cv2.VideoCapture(r'{footage}')\n"
        "ok,frame=cap.read()\n"
        "cap.release()\n"
        f"sys.exit(0 if ok and cv2.imwrite(r'{out_png}',frame) else 1)\n"
    )
    return subprocess.run([str(venv_py), "-c", code]).returncode == 0


def main():
    ap = argparse.ArgumentParser(description="Generate a world with LingBot-World")
    ap.add_argument("--path", required=True, help="lingbot-world install dir")
    ap.add_argument("--model", default="fast", choices=["fast", "cam", "act"])
    ap.add_argument("--image", help="conditioning image")
    ap.add_argument("--footage", help="conditioning video — first frame is extracted")
    ap.add_argument("--prompt", required=True)
    ap.add_argument("--frames", type=int, default=81, help="must be 4n+1")
    ap.add_argument("--size", default="480*832")
    ap.add_argument("--steps", type=int, default=None, help="sampling steps override")
    ap.add_argument("--action-path", help="dir with poses.npy/intrinsics.npy (camera control)")
    ap.add_argument("--action-string", help="keyboard-style action string (act model)")
    ap.add_argument("--out", help="output video file")
    args = ap.parse_args()

    install_path = pathlib.Path(args.path).expanduser()
    venv_py = _venv_python(install_path)
    if not venv_py.exists():
        _fail("not_installed", f"no venv at {install_path} — run `zelpi lingbot install` first")

    manifest = {}
    mpath = install_path / ".zelpi_capabilities.json"
    if mpath.exists():
        try:
            manifest = json.loads(mpath.read_text())
        except Exception:
            pass
    if manifest.get("flash_attn") != True:  # noqa: E712 — "stub" must not pass
        _fail("flash_attn_stub",
              "flash-attn is an import-stub here — LingBot-World's DiT hard-requires the real "
              "kernels, so generation cannot run on this machine (needs Linux + CUDA toolkit, "
              "or a platform with prebuilt flash-attn wheels).")

    ckpt_dir = install_path / "weights" / args.model
    if not ckpt_dir.exists() or not any(ckpt_dir.iterdir()):
        _fail("checkpoint_missing",
              f"no checkpoint at {ckpt_dir} — run `zelpi lingbot pull {args.model}` "
              f"(fast≈75GB, cam/act≈160GB) first")

    if (args.frames - 1) % 4 != 0:
        _fail("bad_frames", f"--frames must be 4n+1 (got {args.frames})")

    image = args.image
    if args.footage:
        footage = pathlib.Path(args.footage)
        if not footage.exists():
            _fail("footage_missing", f"footage not found: {footage}")
        frame_png = install_path / ".zelpi_footage_frame.png"
        print(f"extracting conditioning frame from {footage} …", file=sys.stderr)
        if not _extract_frame(venv_py, footage, frame_png):
            _fail("frame_extract_failed", "could not read a frame from the footage (opencv)")
        image = str(frame_png)
    if not image:
        _fail("no_input", "provide --image or --footage")

    # Low-VRAM knobs from the repo's own CLI. 6-8 GB is still far below what a
    # 27B-class MoE video model wants — warn, add the knobs, try anyway.
    vram_mb = None
    try:
        q = subprocess.run(["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"],
                           capture_output=True, text=True, timeout=15)
        if q.returncode == 0:
            vram_mb = int(q.stdout.strip().split("\n")[0])
    except Exception:
        pass

    script = "generate_fast.py" if args.model == "fast" else "generate.py"
    cmd = [str(venv_py), script,
           "--task", "i2v-A14B",
           "--size", args.size,
           "--ckpt_dir", str(ckpt_dir),
           "--image", str(image),
           "--prompt", args.prompt,
           "--frame_num", str(args.frames)]
    if args.steps:
        cmd += ["--sample_steps", str(args.steps)]
    if args.action_path:
        cmd += ["--action_path", args.action_path]
    if args.action_string:
        cmd += ["--action_string", args.action_string]
    if args.out:
        cmd += ["--save_file", args.out]
    if vram_mb is not None and vram_mb < 24_000:
        print(f"WARNING: {vram_mb} MB VRAM — enabling offload/t5_cpu/dtype-convert; "
              "a 27B-class MoE video model may still not fit.", file=sys.stderr)
        cmd += ["--offload_model", "True", "--t5_cpu", "--convert_model_dtype"]

    print("running:", " ".join(cmd), file=sys.stderr)
    r = subprocess.run(cmd, cwd=install_path)
    if r.returncode != 0:
        _fail("generate_failed", f"{script} exited {r.returncode}")
    print("LINGBOT_GEN_RESULT " + json.dumps({
        "code": "success", "model": args.model, "save_file": args.out or "(repo default naming)",
        "frames": args.frames, "size": args.size,
    }))


if __name__ == "__main__":
    main()