"""Gradio app for the Lyra-2 explorable-scene demo.

Takes a user-uploaded image + caption + trajectory preset and runs Lyra-2's
two-stage pipeline:
  Stage 1 — lyra2_zoomgs_inference:   image -> exploration video (mp4)
  Stage 2 — vipe_da3_gs_recon:        mp4 -> Gaussian splat PLY + cameras.npz

Both stages run in-process against a *resident* Lyra-2 + DA3 (+ MoGe) graph
loaded once at Space startup. The old subprocess path reloaded the DCP
checkpoint on every request (~30 min wasted); this version pays that cost
once, then each request is the ~11 min of actual compute.
"""

from __future__ import annotations

import logging
import os
import shutil
import sys
import threading
import time
import uuid
from dataclasses import dataclass
from pathlib import Path

import gradio as gr
import torch
from PIL import Image

log = logging.getLogger("lyra2-space.app")
logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(message)s")

APP_ROOT = Path("/home/user/app")
LYRA2_DIR = APP_ROOT / "Lyra-2"
RUNS_DIR = Path(os.environ.get("LYRA2_RUNS_DIR", "/tmp/lyra2_runs"))
RUNS_DIR.mkdir(parents=True, exist_ok=True)

# Lyra-2 inference code uses paths like "checkpoints/text_encoder/..." relative
# to LYRA2_DIR. We chdir here once (before importing resident_inference so that
# its module-level torch-state mutations happen with the right CWD) and stay
# here for the lifetime of the process.
if LYRA2_DIR.is_dir():
    os.chdir(LYRA2_DIR)
    if str(LYRA2_DIR) not in sys.path:
        sys.path.insert(0, str(LYRA2_DIR))

# Add the directory containing *this* file to sys.path so resident_inference is
# importable regardless of where Gradio boots us from.
_HERE = Path(__file__).resolve().parent
if str(_HERE) not in sys.path:
    sys.path.insert(0, str(_HERE))

from resident_inference import (  # noqa: E402
    Stage1Resources,
    Stage2Resources,
    build_stage1_args,
    load_stage1_resources,
    load_stage2_resources,
    run_stage1_single,
    run_stage2_single,
)

# Preview video shown while inference is running. Produced by a prior run of
# this exact Space on the pre-filled sample-04 inputs.
PREVIEW_DIR = APP_ROOT / "previews"
SAMPLE04_PREVIEW = PREVIEW_DIR / "sample04_amphitheater.mp4"

INPUT_STEM = "scene"

TRAJECTORY_PRESETS: dict[str, dict[str, float | int]] = {
    "Zoom in + out (default)": {
        "num_frames_zoom_in": 81,
        "num_frames_zoom_out": 241,
        "zoom_in_strength": 0.5,
        "zoom_out_strength": 1.5,
    },
    "Shallow zoom": {
        "num_frames_zoom_in": 81,
        "num_frames_zoom_out": 161,
        "zoom_in_strength": 0.3,
        "zoom_out_strength": 1.0,
    },
    "Wide zoom": {
        "num_frames_zoom_in": 121,
        "num_frames_zoom_out": 321,
        "zoom_in_strength": 0.8,
        "zoom_out_strength": 2.0,
    },
}

SAMPLES_DIR = LYRA2_DIR / "assets" / "samples"
DEFAULT_PRESET = next(iter(TRAJECTORY_PRESETS))

_CURATED_SAMPLE_STEMS = ["04", "05", "09", "12"]


# ---------------------------------------------------------------------------
# Resident model lifecycle
# ---------------------------------------------------------------------------
#
# The Lyra-2 graph is loaded once, then reused across requests. We do this
# lazily on first generate() so that Gradio can bind its HTTP port immediately
# and the Space leaves APP_STARTING quickly — users see the UI during warmup.

_STAGE1_RES: Stage1Resources | None = None
_STAGE2_RES: Stage2Resources | None = None
_WARMUP_LOCK = threading.Lock()
_WARMUP_ERROR: BaseException | None = None
_CONSECUTIVE_FAILURES = 0
_MAX_CONSECUTIVE_FAILURES = 2


def _ensure_resources(progress: gr.Progress | None = None) -> tuple[Stage1Resources, Stage2Resources]:
    """Load Lyra-2 + DA3 (+MoGe) + VIPE once, reuse thereafter."""
    global _STAGE1_RES, _STAGE2_RES, _WARMUP_ERROR
    if _STAGE1_RES is not None and _STAGE2_RES is not None:
        return _STAGE1_RES, _STAGE2_RES

    with _WARMUP_LOCK:
        if _STAGE1_RES is not None and _STAGE2_RES is not None:
            return _STAGE1_RES, _STAGE2_RES
        if _WARMUP_ERROR is not None:
            raise gr.Error(
                f"Model warmup failed on a previous attempt: {_WARMUP_ERROR}. "
                "Check Space logs."
            )

        t_start = time.monotonic()
        if progress is not None:
            progress(0.01, desc="Loading Lyra-2 (one-time, ~30 min on cold A100)...")
        log.info("[warmup] starting stage-1 resource load")
        try:
            args = build_stage1_args(
                checkpoint_dir="checkpoints/model",
                experiment="lyra2",
                use_dmd=True,
            )
            stage1 = load_stage1_resources(args)
            log.info(
                "[warmup] stage-1 done in %.1fs; loading stage-2",
                time.monotonic() - t_start,
            )
            if progress is not None:
                progress(0.04, desc="Loading VIPE (one-time)...")
            stage2 = load_stage2_resources(da3_from_stage1=stage1.da3_model)
            _STAGE1_RES = stage1
            _STAGE2_RES = stage2
            log.info(
                "[warmup] full warmup complete in %.1fs",
                time.monotonic() - t_start,
            )
        except BaseException as e:  # noqa: BLE001 - cache so the next request errors fast
            _WARMUP_ERROR = e
            raise

    assert _STAGE1_RES is not None and _STAGE2_RES is not None
    return _STAGE1_RES, _STAGE2_RES


def _load_bundled_examples() -> list[list]:
    examples: list[list] = []
    for stem in _CURATED_SAMPLE_STEMS:
        img = SAMPLES_DIR / f"{stem}.png"
        txt = SAMPLES_DIR / f"{stem}.txt"
        if not (img.is_file() and txt.is_file()):
            log.info("Skipping example %s: files missing at %s", stem, SAMPLES_DIR)
            continue
        caption = txt.read_text(encoding="utf-8").strip()
        examples.append([str(img), caption, DEFAULT_PRESET, True])
    return examples


@dataclass
class GenerationOutputs:
    video_path: Path
    ply_path: Path


def _write_input_image(run_dir: Path, image: Image.Image) -> Path:
    inputs_dir = run_dir / "inputs"
    inputs_dir.mkdir(parents=True, exist_ok=True)
    img_path = inputs_dir / f"{INPUT_STEM}.png"
    image.convert("RGB").save(img_path, format="PNG")
    return img_path


def _stage1_video(
    run_dir: Path,
    input_image: Path,
    prompt: str,
    preset_params: dict,
    res1: Stage1Resources,
    progress: gr.Progress,
) -> Path:
    progress(0.05, desc="Stage 1/2 — generating exploration video (~8 min)")
    t0 = time.monotonic()
    video_path = run_stage1_single(
        res1,
        image_path=input_image,
        prompt=prompt,
        preset_params=preset_params,
        output_path=run_dir / "zoomgs",
    )
    log.info("[stage1] done in %.1fs -> %s", time.monotonic() - t0, video_path)
    if not video_path.exists():
        raise gr.Error(f"Stage 1 did not produce {video_path}.")
    return video_path


def _stage2_splat(
    run_dir: Path,
    video_path: Path,
    res1: Stage1Resources,
    res2: Stage2Resources,
    progress: gr.Progress,
) -> Path:
    progress(0.75, desc="Stage 2/2 — reconstructing 3D Gaussian splat (~3 min)")
    t0 = time.monotonic()

    # Offload Lyra-2 diffusion net to CPU so DA3 GS recon has VRAM headroom.
    # First-run measurement on A100-80GB: with stage-1 resident, DA3's DINOv2
    # backbone OOMs trying to allocate 1.47 GiB (78.98/79.25 GiB already used).
    # DA3 stays on GPU since it's reused across both stages.
    res1.model.net.cpu()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    try:
        ply_path = run_stage2_single(
            res2,
            video_path=video_path,
            output_dir=run_dir / "recon",
        )
    finally:
        # Restore the diffusion net to GPU for the next stage-1 request.
        res1.model.net.to(device=res1.desired_device, dtype=res1.desired_dtype)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    log.info("[stage2] done in %.1fs -> %s", time.monotonic() - t0, ply_path)
    if not ply_path.exists():
        raise gr.Error(f"Stage 2 did not produce {ply_path}.")
    return ply_path


def _new_run_dir() -> Path:
    run_dir = RUNS_DIR / uuid.uuid4().hex[:12]
    run_dir.mkdir(parents=True, exist_ok=True)
    return run_dir


def _prune_old_runs(keep_last: int = 3) -> None:
    runs = sorted(
        (p for p in RUNS_DIR.iterdir() if p.is_dir()),
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )
    for stale in runs[keep_last:]:
        shutil.rmtree(stale, ignore_errors=True)


def generate(
    image: Image.Image | None,
    prompt: str,
    preset_name: str,
    use_dmd: bool,
    progress: gr.Progress = gr.Progress(track_tqdm=False),
) -> tuple[str, str]:
    global _CONSECUTIVE_FAILURES
    if image is None:
        raise gr.Error("Please upload an image.")
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a caption describing your scene.")
    preset_params = TRAJECTORY_PRESETS.get(preset_name)
    if preset_params is None:
        raise gr.Error(f"Unknown trajectory preset: {preset_name}")

    res1, res2 = _ensure_resources(progress=progress)

    _prune_old_runs()
    run_dir = _new_run_dir()
    log.info("run_dir=%s preset=%s", run_dir, preset_name)

    input_image = _write_input_image(run_dir, image)

    try:
        video_path = _stage1_video(run_dir, input_image, prompt.strip(), preset_params, res1, progress)
        ply_path = _stage2_splat(run_dir, video_path, res1, res2, progress)
    except gr.Error:
        _CONSECUTIVE_FAILURES += 1
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        _maybe_hard_exit()
        raise
    except BaseException as e:  # noqa: BLE001 — convert to gr.Error with counter bookkeeping
        _CONSECUTIVE_FAILURES += 1
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        log.exception("[generate] inference failed")
        _maybe_hard_exit()
        raise gr.Error(f"Inference failed: {type(e).__name__}: {e}") from e

    _CONSECUTIVE_FAILURES = 0
    progress(1.0, desc="Complete")
    return str(video_path), str(ply_path)


def _maybe_hard_exit() -> None:
    """After N consecutive failures, take the container down so HF respawns a
    fresh Space. A resident torch graph can corrupt under OOM / cuDNN faults
    and a respawn is cheaper than guessing which internal state is bad.
    """
    if _CONSECUTIVE_FAILURES >= _MAX_CONSECUTIVE_FAILURES:
        log.error(
            "[generate] %d consecutive failures — exiting so HF respawns the Space",
            _CONSECUTIVE_FAILURES,
        )
        os._exit(1)


def _default_sample() -> tuple[str | None, str]:
    first = _CURATED_SAMPLE_STEMS[0]
    img = SAMPLES_DIR / f"{first}.png"
    txt = SAMPLES_DIR / f"{first}.txt"
    if img.is_file() and txt.is_file():
        return str(img), txt.read_text(encoding="utf-8").strip()
    return None, ""


def _build_ui() -> gr.Blocks:
    bundled_examples = _load_bundled_examples()
    default_image, default_caption = _default_sample()
    with gr.Blocks(title="Lyra-2 Explorable Scene Demo") as ui:
        gr.Markdown(
            """
            # Lyra-2 — image to explorable 3D scene

            Upload an image + caption. Get back an exploration video and a
            Gaussian-splat `.ply`. Powered by
            [nvidia/Lyra-2.0](https://huggingface.co/nvidia/Lyra-2.0).

            **Runtime: ~13 min per request on A100 80GB.** The Space pays a
            one-time ~60 min cold-boot before the UI shows up; by the time
            you can click Generate, the model is already resident. Queue is
            serial — one request at a time.
            """
        )
        with gr.Row():
            with gr.Column(scale=1):
                image_in = gr.Image(
                    type="pil",
                    label="Input image",
                    sources=["upload", "clipboard"],
                    value=default_image,
                )
                prompt_in = gr.Textbox(
                    label="Caption",
                    placeholder="Describe the scene...",
                    lines=3,
                    max_lines=5,
                    value=default_caption,
                )
                preset_in = gr.Dropdown(
                    choices=list(TRAJECTORY_PRESETS.keys()),
                    value=list(TRAJECTORY_PRESETS.keys())[0],
                    label="Trajectory preset",
                )
                dmd_in = gr.Checkbox(
                    value=True,
                    label="Fast mode (DMD distillation)",
                    info="Locked on for this demo. Disabling it would push inference past 30 min.",
                    interactive=False,
                )
                run_btn = gr.Button("Generate scene", variant="primary")
            with gr.Column(scale=1):
                if SAMPLE04_PREVIEW.exists():
                    gr.Markdown(
                        "### Sample output (from a prior run on the pre-filled input)"
                    )
                    gr.Video(
                        value=str(SAMPLE04_PREVIEW),
                        label="Preview — this exact input produced this video on a previous run",
                        autoplay=True,
                        loop=True,
                        interactive=False,
                    )
                    gr.Markdown("---")
                    gr.Markdown("### Your result")
                video_out = gr.Video(label="Exploration video", autoplay=False, loop=True)
                ply_out = gr.File(
                    label="Gaussian splat (.ply) — download and open in any GS viewer",
                    file_count="single",
                )
                gr.Markdown(
                    """
                    ### Walk the scene on macOS — *coming soon*
                    The downloadable `.ply` is a standard binary Gaussian splat and
                    works with any GS viewer today. A dedicated macOS viewer for a
                    walkable first-person experience (WASD + mouse-look) is in the
                    works — stay tuned.
                    """
                )
        if bundled_examples:
            gr.Examples(
                examples=bundled_examples,
                inputs=[image_in, prompt_in, preset_in, dmd_in],
                label="Try a bundled example — click to populate all inputs",
                examples_per_page=4,
            )
        run_btn.click(
            fn=generate,
            inputs=[image_in, prompt_in, preset_in, dmd_in],
            outputs=[video_out, ply_out],
            concurrency_limit=1,  # single GPU + resident non-reentrant model
        )
    return ui


if __name__ == "__main__":
    # Warmup BEFORE Gradio binds — HF's gcTimeout (15 min) starts as soon as the
    # HTTP port is up, and it kills the container if no requests flow. Our silent
    # 60-min DCP load looks "idle" to that GC. Keeping the Space in APP_STARTING
    # through warmup lets startup_duration_timeout (2h in README) cover us.
    log.info("[main] starting synchronous model warmup (expected ~60 min)")
    _ensure_resources(progress=None)
    log.info("[main] warmup done; launching Gradio")
    ui = _build_ui()
    ui.queue(max_size=8).launch(
        server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),
        server_port=int(os.environ.get("GRADIO_SERVER_PORT", "7860")),
    )