"""TripoSplat – gradio.Server with custom frontend.
Usage: python app.py
"""
import base64
import os
import subprocess
import tempfile
import time
from pathlib import Path
from uuid import uuid4

import spaces
import torch
from PIL import Image
from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
from gradio import Server
from gradio.data_classes import FileData

from triposplat import TripoSplatPipeline
import example_inputs_b64 as _b64

# ----------------------------------------------------------------------------
# Download checkpoints from HuggingFace Hub (VAST-AI/TripoSplat)
# ----------------------------------------------------------------------------

subprocess.run(
    [
        "hf", "download",
        "VAST-AI/TripoSplat",
        "--local-dir", "ckpts"
    ],
    check=True,
)

# ----------------------------------------------------------------------------
# Pipeline (loaded once at startup)
# ----------------------------------------------------------------------------

PIPE = TripoSplatPipeline(
    ckpt_path              = "ckpts/diffusion_models/triposplat_fp16.safetensors",
    decoder_path           = "ckpts/vae/triposplat_vae_decoder_fp16.safetensors",
    dinov3_path            = "ckpts/clip_vision/dino_v3_vit_h.safetensors",
    flux2_vae_encoder_path = "ckpts/vae/flux2-vae.safetensors",
    rmbg_path              = "ckpts/background_removal/birefnet.safetensors",
    device                 = "cuda",
)

OUT_ROOT = Path("gradio_outputs").resolve()
OUT_ROOT.mkdir(parents=True, exist_ok=True)

# Decode example images from base64 into a persistent temp directory so that
# the custom frontend can serve them via FastAPI routes.
_EXAMPLES_TMPDIR = tempfile.mkdtemp(prefix="triposplat_examples_")


def _write_example(varname: str, filename: str) -> str:
    path = Path(_EXAMPLES_TMPDIR) / filename
    path.write_bytes(base64.b64decode(getattr(_b64, varname)))
    return str(path)


EXAMPLES = [
    {"name": "Creature Butterfly",  "file": _write_example("CREATURE_BUTTERFLY",   "creature_butterfly.webp")},
    {"name": "Building Stone House","file": _write_example("BUILDING_STONE_HOUSE", "building_stone_house.webp")},
    {"name": "Vehicle Pirate Ship", "file": _write_example("VEHICLE_PIRATE_SHIP",  "vehicle_pirate_ship.webp")},
    {"name": "Plant Water Lily",    "file": _write_example("PLANT_WATER_LILY",     "plant_water_lily.webp")},
]

# ----------------------------------------------------------------------------
# gradio.Server
# ----------------------------------------------------------------------------

app = Server()


# ---- Static pages ----------------------------------------------------------

@app.get("/")
async def homepage():
    """Serve the custom frontend."""
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return HTMLResponse(f.read())


@app.get("/viewer")
async def viewer_page():
    """Serve the Spark.js 3D viewer (loaded inside an iframe)."""
    viewer_path = os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "static", "viewer", "viewer.html",
    )
    with open(viewer_path, "r", encoding="utf-8") as f:
        return HTMLResponse(f.read())


# ---- Example images --------------------------------------------------------

@app.get("/api/examples")
async def get_examples():
    """Return a JSON list of example images the frontend can display."""
    return JSONResponse([
        {"name": ex["name"], "url": f"/api/example/{i}"}
        for i, ex in enumerate(EXAMPLES)
    ])


@app.get("/api/example/{idx}")
async def get_example(idx: int):
    """Serve an individual example image by index."""
    if 0 <= idx < len(EXAMPLES):
        return FileResponse(EXAMPLES[idx]["file"], media_type="image/webp")
    return JSONResponse({"error": "not found"}, status_code=404)


# ----------------------------------------------------------------------------
# GPU pipeline helper
# ----------------------------------------------------------------------------

@spaces.GPU
def _run_pipeline(pil_image, seed, steps, guidance_scale, num_gaussians,
                  out_dir, output_format):
    """Run the full pipeline (preprocess → encode → sample → decode → save)
    in a single GPU acquisition.

    All file I/O happens here so the unpicklable Gaussian object never
    crosses the ZeroGPU multiprocessing boundary.
    """
    t0 = time.time()
    prepared = PIPE.preprocess_image(pil_image)
    gen = torch.Generator(device=PIPE._device).manual_seed(int(seed))
    cond = PIPE.encode_image(prepared, generator=gen)
    out = PIPE.sample_latent(
        cond,
        steps=int(steps),
        guidance_scale=float(guidance_scale),
        generator=gen,
        show_progress=True,
    )
    gaussian = PIPE.decode_latent(out["latent"], num_gaussians=int(num_gaussians))
    gen_dt = time.time() - t0

    # Save preprocessed image
    prep_path = out_dir / "preprocessed.png"
    prepared.save(str(prep_path))

    # Save PLY (always needed for the viewer)
    ply_path = out_dir / "splat.ply"
    gaussian.save_ply(str(ply_path))

    # Save in the requested download format
    fmt = output_format.lower()
    if fmt == "splat":
        download_path = out_dir / "splat.splat"
        gaussian.save_splat(str(download_path))
    else:
        download_path = ply_path

    n_gaussians = gaussian.get_xyz.shape[0]

    # Return only picklable primitives / paths
    return str(prep_path), str(ply_path), str(download_path), n_gaussians, gen_dt


# ----------------------------------------------------------------------------
# Main API endpoint  (queued via Gradio's engine)
# ----------------------------------------------------------------------------

@app.api()
def generate(
    image: FileData,
    seed: int = 42,
    steps: int = 20,
    guidance_scale: float = 3.0,
    num_gaussians: int = 262144,
    output_format: str = "ply",
) -> tuple[FileData, FileData, FileData, str]:
    """Generate 3D Gaussians from an input image.

    Returns (preprocessed_image, ply_file, download_file, info_string).
    The frontend receives these as result.data[0..3].
    """
    pil_image = Image.open(image["path"]).convert("RGBA")

    out_dir = OUT_ROOT / uuid4().hex[:12]
    out_dir.mkdir(parents=True, exist_ok=True)

    prep_path, ply_path, download_path, n_gaussians, gen_dt = _run_pipeline(
        pil_image, seed, steps, guidance_scale, num_gaussians,
        out_dir, output_format,
    )

    info = (
        f"{n_gaussians:,} gaussians  ·  "
        f"generation: {gen_dt:.1f}s  ·  saved: {Path(download_path).name}"
    )

    return (
        FileData(path=prep_path),
        FileData(path=ply_path),
        FileData(path=download_path),
        info,
    )


# ----------------------------------------------------------------------------
# Launch
# ----------------------------------------------------------------------------

if __name__ == "__main__":
    app.launch(show_error=True)