"""
Hunyuan3D-2 — Shape-only HuggingFace Space
Uses Hunyuan3D-2mini-Turbo (0.6 B, step-distilled) for fast shape generation
within standard ZeroGPU quota. No texture pipeline — mesh only.
Background removal is skipped — designed for clean character art.
"""

import os
import tempfile

import gradio as gr
import spaces          # ZeroGPU decorator
import torch
from PIL import Image

# ---------------------------------------------------------------------------
# Lazy global pipeline — loaded once on first GPU call
# ---------------------------------------------------------------------------
_pipeline = None


def _get_pipeline():
    """Load the shape pipeline once and cache it globally."""
    global _pipeline
    if _pipeline is None:
        from hy3dgen.shapegen import Hunyuan3DDiTFlowMatchingPipeline

        _pipeline = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(
            "tencent/Hunyuan3D-2mini",
            subfolder="hunyuan3d-dit-v2-mini-turbo",
            use_safetensors=True,
            torch_dtype=torch.float16,
        )
    return _pipeline


# ---------------------------------------------------------------------------
# Simple image pre-processing — resize only, no background removal
# ---------------------------------------------------------------------------
def preprocess_image(pil_image: Image.Image) -> Image.Image:
    """Resize to 512x512 RGB — the model's native conditioning resolution."""
    return pil_image.convert("RGB").resize((512, 512), Image.LANCZOS)


# ---------------------------------------------------------------------------
# Core generation — wrapped in @spaces.GPU for ZeroGPU
# ---------------------------------------------------------------------------
@spaces.GPU(duration=60)
def generate_shape(image: Image.Image, seed: int, steps: int, octree_res: int):
    """
    Run Hunyuan3D-DiT shape generation and return a GLB file path.

    NOTE: pipeline.to("cuda") must NOT be reassigned — some custom pipelines
    return None from .to(), which would make the pipeline uncallable.
    """
    pipeline = _get_pipeline()
    pipeline.to("cuda")   # move in-place; do not reassign the return value

    generator = torch.Generator(device="cuda").manual_seed(seed)

    meshes = pipeline(
        image=image,
        num_inference_steps=steps,
        octree_resolution=octree_res,
        num_chunks=8000,
        generator=generator,
        output_type="trimesh",
    )
    mesh = meshes[0]

    tmp_dir = tempfile.mkdtemp()
    out_path = os.path.join(tmp_dir, "shape.glb")
    mesh.export(out_path)
    return out_path


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
def run(image, seed, steps, octree_res, progress=gr.Progress(track_tqdm=True)):
    if image is None:
        raise gr.Error("Please upload an image first.")

    progress(0.1, desc="Preprocessing image ...")
    pil = image if isinstance(image, Image.Image) else Image.fromarray(image)
    processed = preprocess_image(pil)

    progress(0.3, desc="Running shape diffusion ...")
    glb_path = generate_shape(processed, int(seed), int(steps), int(octree_res))

    progress(1.0, desc="Done!")
    return glb_path, processed, glb_path


with gr.Blocks(title="Hunyuan3D-2 Shape Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Hunyuan3D-2 Shape Generator
        Upload character art to generate an **untextured 3-D mesh** using
        [Hunyuan3D-2mini-Turbo](https://huggingface.co/tencent/Hunyuan3D-2mini).
        Shape only - no texture - stays well within the ZeroGPU free quota.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="Input Image",
                type="pil",
                sources=["upload", "clipboard"],
                height=340,
            )

            with gr.Accordion("Advanced settings", open=False):
                seed = gr.Slider(
                    label="Seed",
                    minimum=0, maximum=2**31 - 1,
                    value=42, step=1,
                )
                steps = gr.Slider(
                    label="Diffusion steps",
                    minimum=5, maximum=50,
                    value=5, step=1,
                    info="5-15 works well with the turbo model.",
                )
                octree_res = gr.Slider(
                    label="Octree resolution",
                    minimum=128, maximum=512,
                    value=192, step=64,
                    info="Higher = finer mesh detail but more VRAM & time.",
                )

            generate_btn = gr.Button("Generate Shape", variant="primary")

        with gr.Column(scale=1):
            preview_img = gr.Image(
                label="Image sent to model (512x512)",
                type="pil",
                interactive=False,
                height=200,
            )
            output_3d = gr.Model3D(
                label="3-D Shape (GLB)",
                height=400,
                clear_color=[0.9, 0.9, 0.9, 1.0],
            )
            download_file = gr.File(label="Download GLB")

    gr.Markdown(
        """
        ---
        **Tips**
        - Works best with clean character art on a plain or transparent background.
        - Lower octree resolution (128-192) is faster and still looks great for most art.
        - Model: Hunyuan3D-DiT-v2-mini-Turbo - 0.6B parameters, step-distilled.
        """
    )

    generate_btn.click(
        fn=run,
        inputs=[input_image, seed, steps, octree_res],
        outputs=[output_3d, preview_img, download_file],
    )


if __name__ == "__main__":
    demo.queue(max_size=5).launch()