import gradio as gr
import torch
import spaces

import json
import os
import random
import subprocess
import sys
import uuid
from pathlib import Path

os.environ.setdefault("SPCONV_ALGO", "native")
os.environ.setdefault("ATTN_BACKEND", "flash_attn")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("DVD_MODEL_REPO", "Zhengrui/dvd")

MAX_SEED = 2**31 - 1
TMP_DIR = Path(__file__).resolve().parent / "tmp" / "dvd_image_min"
TMP_DIR.mkdir(parents=True, exist_ok=True)


def log_event(message: str):
    print(f"[DVD Image Min] {message}", flush=True)


def worker_path(name: str) -> str:
    path = TMP_DIR / f"worker-{uuid.uuid4().hex}"
    path.mkdir(parents=True, exist_ok=True)
    return str(path / name)


@spaces.GPU(duration=30)
def zero_gpu_smoke_test():
    log_event("zero_gpu_smoke_test start")
    if not torch.cuda.is_available():
        log_event("cuda unavailable")
        return "CUDA unavailable"
    value = torch.ones((1,), device="cuda").sum().item()
    name = torch.cuda.get_device_name(0)
    log_event(f"zero_gpu_smoke_test done device={name} value={value}")
    return f"OK: {name}, value={value}"


@spaces.GPU(duration=120)
def generate_voxels(image, seed, randomize_seed, preprocess_image, dvd_steps, progress=gr.Progress(track_tqdm=True)):
    progress(0.01, desc="Starting ZeroGPU callback")
    log_event(f"generate_voxels start seed={seed} randomize={randomize_seed} steps={dvd_steps}")
    if image is None:
        raise gr.Error("Please provide an image.")
    seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    image_path = worker_path("input.png")
    mesh_path = worker_path("generated_voxels.glb")
    npy_path = worker_path("generated_voxel64_coords.npy")
    config_path = worker_path("generate_config.json")
    image.save(image_path)
    with open(config_path, "w") as f:
        json.dump(
            {
                "image_path": image_path,
                "mesh_path": mesh_path,
                "npy_path": npy_path,
                "seed": int(seed),
                "preprocess_image": bool(preprocess_image),
                "dvd_steps": int(dvd_steps),
                "dvd_cfg_mode": "Default schedule",
                "dvd_cfg_constant": 0.7,
                "dvd_cfg_early": 0.4,
                "dvd_cfg_late": 0.7,
                "dvd_cfg_split": 0.5,
            },
            f,
        )
    progress(0.08, desc="Running isolated DVD worker")
    env = os.environ.copy()
    env.setdefault("PYTHONUNBUFFERED", "1")
    cmd = [sys.executable, str(Path(__file__).resolve().parent / "space_image_worker.py"), "generate", config_path]
    proc = subprocess.run(cmd, env=env, text=True, capture_output=True)
    if proc.stdout:
        print(proc.stdout, flush=True)
    if proc.stderr:
        print(proc.stderr, flush=True)
    if proc.returncode != 0:
        raise gr.Error(f"DVD worker failed with exit code {proc.returncode}. Check Space logs.")
    log_event(f"generate_voxels done seed={seed}")
    return mesh_path, npy_path, int(seed), f"Done. seed={seed}"


with gr.Blocks(title="DVD Image", fill_width=True) as demo:
    with gr.Row():
        smoke_btn = gr.Button("ZeroGPU Smoke Test")
        smoke_out = gr.Textbox(label="ZeroGPU Status", interactive=False)
    smoke_btn.click(zero_gpu_smoke_test, outputs=smoke_out)

    with gr.Row(equal_height=False):
        with gr.Column():
            image = gr.Image(label="Input Image", format="png", image_mode="RGBA", type="pil", height=320)
            seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
            randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
            preprocess_image = gr.Checkbox(value=True, label="Preprocess image")
            dvd_steps = gr.Slider(1, 512, value=256, step=1, label="DVD steps")
            gen_btn = gr.Button("Generate DVD Voxels", variant="primary")
        with gr.Column():
            voxel_view = gr.Model3D(label="Generated / Cubified Voxels", height=360, camera_position=(-180, 90, 3))
            npy_download = gr.DownloadButton(label="Download Voxel Coords (.npy)", interactive=False)
            status = gr.Textbox(label="Status", interactive=False)

    gen_btn.click(
        generate_voxels,
        inputs=[image, seed, randomize_seed, preprocess_image, dvd_steps],
        outputs=[voxel_view, npy_download, seed, status],
    ).then(lambda: gr.DownloadButton(interactive=True), outputs=[npy_download])


if __name__ == "__main__":
    demo.queue().launch(show_api=False, show_error=True, ssr_mode=False)