#!/usr/bin/env python3
"""
lingbot_install.py — install LingBot-World (github.com/robbyant/lingbot-world),
Ant Group's open interactive world model (Wan2.2-class video generation:
image/footage + text + camera/action control → explorable video world),
into an isolated venv. Follows the exact pattern of hyworld_install.py
(disk-safe pip env, capabilities manifest, JSON sentinel result) and reuses
its helpers directly.

Checkpoints are NOT downloaded here — they are 74 GB (fast) to 160 GB
(base-act) and get their own explicit `zelpi lingbot pull` step with a disk
preflight, so `lingbot install` stays a minutes-not-hours operation.

Honest capability notes recorded in the manifest:
  - torch_cuda: whether the venv's torch can actually see a GPU
  - flash_attn: LingBot-World's generate path hard-requires flash-attn.
    There are no official Windows wheels; without nvcc it CANNOT be built,
    the stub only satisfies imports — generation will not run without the
    real kernels, and the manifest says so rather than pretending.
"""
from __future__ import annotations

import argparse
import json
import pathlib
import shutil
import subprocess
import sys

# Sibling-module reuse: when run as a script, sys.path[0] is this directory.
import hyworld_install as base

log = base.log

REPO_URL = "https://github.com/robbyant/lingbot-world"

CHECKPOINTS = {
    "fast": {"hf": "robbyant/lingbot-world-fast", "size_gb": 75},
    "cam": {"hf": "robbyant/lingbot-world-base-cam", "size_gb": 160},
    "act": {"hf": "robbyant/lingbot-world-base-act", "size_gb": 161},
}


def main():
    ap = argparse.ArgumentParser(description="Install LingBot-World into an isolated venv")
    ap.add_argument("--path", default=str(pathlib.Path.home() / ".zelpi" / "lingbot-world"))
    ap.add_argument("--cuda", default="cu126", help="CUDA wheel channel when a GPU is present")
    ap.add_argument("--no-gpu", action="store_true", help="skip CUDA torch even if a GPU is detected")
    args = ap.parse_args()

    import logging
    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(name)s  %(levelname)s  %(message)s", datefmt="%H:%M:%S")

    install_path = pathlib.Path(args.path).expanduser()
    manifest_path = install_path / ".zelpi_capabilities.json"
    capabilities = {"torch_cuda": False, "flash_attn": False, "python": None, "path": str(install_path), "checkpoints": {}}
    if manifest_path.exists():
        try:
            capabilities.update(json.loads(manifest_path.read_text(encoding="utf-8")))
        except Exception:
            pass
    result = {"code": "success", "capabilities": capabilities}

    log.info("─── PREFLIGHT ───")
    base._check_disk_space(install_path, required_gb=15)
    vram = base._check_gpu_vram()

    log.info("─── CLONE/PULL REPO ───")
    if not base._clone_or_pull(REPO_URL, install_path):
        result["code"] = "clone_failed"
        print("LINGBOT_INSTALL_RESULT " + json.dumps(result))
        sys.exit(1)

    log.info("─── CREATE VENV ───")
    if not base._create_venv(install_path):
        result["code"] = "venv_failed"
        print("LINGBOT_INSTALL_RESULT " + json.dumps(result))
        sys.exit(1)

    venv_python = base._get_venv_python(install_path)
    capabilities["python"] = str(venv_python)
    pip_env = base._pip_env(install_path)

    log.info("─── UPGRADE PIP/SETUPTOOLS ───")
    subprocess.run([str(venv_python), "-m", "pip", "install", "--upgrade", "pip", "setuptools", "wheel"], env=pip_env, check=False)

    # torch first, and from the CUDA index when a GPU is present — LingBot-World
    # requires torch>=2.4, and installing it FIRST from the right index avoids
    # requirements.txt resolving a CPU-only build that we'd have to swap later.
    log.info("─── TORCH ───")
    use_gpu = vram is not None and not args.no_gpu
    torch_args = [str(venv_python), "-m", "pip", "install", "--timeout", "120", "torch", "torchvision"]
    if use_gpu:
        torch_args += ["--index-url", f"https://download.pytorch.org/whl/{args.cuda}"]
    try:
        subprocess.run(torch_args, check=True, env=pip_env)
    except subprocess.CalledProcessError:
        result["code"] = "torch_failed"
        print("LINGBOT_INSTALL_RESULT " + json.dumps(result))
        sys.exit(1)
    r = subprocess.run([str(venv_python), "-c", "import torch;print(torch.cuda.is_available())"],
                       capture_output=True, text=True, env=pip_env)
    capabilities["torch_cuda"] = args.cuda if r.stdout.strip().endswith("True") else False
    log.info("torch CUDA available: %s", r.stdout.strip())

    # flash_attn sits INSIDE the repo's requirements.txt — on any machine
    # without a CUDA toolchain its source build fails and would take the
    # whole `-r requirements.txt` install down with it. Install everything
    # else first from a filtered copy, then handle flash-attn on its own.
    log.info("─── REQUIREMENTS (flash_attn handled separately) ───")
    reqs = (install_path / "requirements.txt").read_text(encoding="utf-8").splitlines()
    filtered = [l for l in reqs if not l.strip().lower().replace("-", "_").startswith("flash_attn")]
    filtered_name = ".zelpi_requirements_no_flash.txt"
    (install_path / filtered_name).write_text("\n".join(filtered) + "\n", encoding="utf-8")
    if not base._pip_install(venv_python, install_path, filtered_name, pip_env):
        result["code"] = "pip_failed"
        print("LINGBOT_INSTALL_RESULT " + json.dumps(result))
        sys.exit(1)

    # hf_transfer: makes the 74-160 GB `zelpi lingbot pull` checkpoint downloads
    # survive flaky CDN connections (huggingface_hub picks it up automatically
    # when HF_HUB_ENABLE_HF_TRANSFER=1; huggingface_hub itself is already here
    # as a transformers/diffusers dependency).
    base._pip_install_package(venv_python, "hf_transfer", pip_env)

    # flash-attn: hard-required by the generation path (the repo's own install
    # step is `pip install flash-attn --no-build-isolation`). Only buildable
    # with a CUDA toolchain; otherwise install the import-stub and record the
    # honest consequence: generation will NOT run until real kernels exist.
    log.info("─── FLASH-ATTN ───")
    if shutil.which("nvcc"):
        try:
            subprocess.run([str(venv_python), "-m", "pip", "install", "--timeout", "120",
                            "flash-attn", "--no-build-isolation"], check=True, env=pip_env)
            capabilities["flash_attn"] = True
            log.info("✓ flash-attn built")
        except subprocess.CalledProcessError:
            log.warning("✗ flash-attn build failed — generation will not run without it")
            if base._install_flash_attn_stub(venv_python, pip_env):
                capabilities["flash_attn"] = "stub"
    else:
        log.warning("no nvcc on PATH — flash-attn cannot be built here (no official Windows wheels); "
                    "installing import-stub. LingBot-World GENERATION WILL NOT RUN on this machine "
                    "until flash-attn is really installed (Linux + CUDA toolkit).")
        if base._install_flash_attn_stub(venv_python, pip_env):
            capabilities["flash_attn"] = "stub"

    manifest_path.write_text(json.dumps(capabilities, indent=2), encoding="utf-8")
    log.info("─── DONE ───")
    result["capabilities"] = capabilities
    print("LINGBOT_INSTALL_RESULT " + json.dumps(result))


if __name__ == "__main__":
    main()