import spaces  # must come before torch
import os
import random
import tempfile
import uuid

import gradio as gr
import numpy as np
import torch
from diffusers import ControlNetModel, EulerAncestralDiscreteScheduler, StableDiffusionControlNetPipeline
from PIL import Image

import uvpos

BASE_MODEL = "stable-diffusion-v1-5/stable-diffusion-v1-5"
CONTROLNET = "GeorgeQi/Paint3d_UVPos_Control"
RES = 1024  # Paint3D generates UV maps at 1024x1024
MAX_SEED = 2**31 - 1
DEFAULT_NEG = (
    "blur, low quality, noisy image, over-exposed, strong light, bright light, "
    "shadow, darkness, silhouette, specular highlights"
)

controlnet = ControlNetModel.from_pretrained(CONTROLNET, torch_dtype=torch.float16)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    BASE_MODEL, controlnet=controlnet, torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False
)
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.to("cuda")

OUT_DIR = os.path.join(tempfile.gettempdir(), "paint3d_out")
os.makedirs(OUT_DIR, exist_ok=True)


def make_uv_position(mesh_file: str):
    """Rasterize the mesh's UV position map (the ControlNet condition).

    Args:
        mesh_file: Path to a .glb/.gltf/.obj/.ply mesh. Meshes without UVs are unwrapped with xatlas.

    Returns:
        The 1024x1024 UV position map image.
    """
    if not mesh_file:
        raise gr.Error("Upload a mesh first.")
    _, _, mesh = uvpos.load_mesh(mesh_file)
    pos, _ = uvpos.rasterize_uv_position(mesh, RES)
    return Image.fromarray((pos * 255).astype(np.uint8))


@spaces.GPU(duration=60)
def _generate(cond: Image.Image, prompt: str, negative_prompt: str, steps: int, guidance: float,
              control_scale: float, seed: int) -> Image.Image:
    g = torch.Generator("cuda").manual_seed(int(seed))
    return pipe(
        prompt, negative_prompt=negative_prompt, image=cond, width=RES, height=RES,
        num_inference_steps=int(steps), guidance_scale=float(guidance),
        controlnet_conditioning_scale=float(control_scale), generator=g,
    ).images[0]


def texture_mesh(
    mesh_file: str,
    prompt: str,
    negative_prompt: str = DEFAULT_NEG,
    steps: int = 30,
    guidance: float = 7.0,
    control_scale: float = 1.0,
    keep_normal_map: bool = True,
    seed: int = 0,
    randomize_seed: bool = True,
):
    """Generate a lighting-less UV texture for a 3D mesh from a text prompt with Paint3D's UV-position ControlNet.

    Args:
        mesh_file: Path to a .glb/.gltf/.obj/.ply mesh (UVs are used if present, otherwise xatlas-unwrapped).
        prompt: Text description of the desired materials/appearance.
        negative_prompt: What to avoid in the texture.
        steps: Number of diffusion steps.
        guidance: Classifier-free guidance scale.
        control_scale: Strength of the UV-position ControlNet conditioning.
        keep_normal_map: Keep the input GLB's normal map on the output material if it has one.
        seed: Random seed.
        randomize_seed: Pick a random seed instead of `seed`.

    Returns:
        Textured GLB path, generated UV texture, UV position map, and the seed used.
    """
    if not mesh_file:
        raise gr.Error("Upload a mesh first.")
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    scene, name, mesh = uvpos.load_mesh(mesh_file)
    pos, mask = uvpos.rasterize_uv_position(mesh, RES)
    cond = Image.fromarray((pos * 255).astype(np.uint8))

    full_prompt = prompt.strip()
    if not full_prompt.lower().startswith("uv"):
        full_prompt = f"UV map, {full_prompt}, high quality"
    tex = _generate(cond, full_prompt, negative_prompt, steps, guidance, control_scale, seed)

    tex_np = uvpos.dilate_texture(np.asarray(tex.convert("RGB")), mask)
    tex = Image.fromarray(tex_np)

    uid = uuid.uuid4().hex[:8]
    glb_path = uvpos.export_textured(scene, name, mesh, tex, keep_normal_map, os.path.join(OUT_DIR, f"textured_{uid}.glb"))
    tex_path = os.path.join(OUT_DIR, f"texture_{uid}.png")
    tex.save(tex_path)
    return glb_path, tex_path, cond, seed


EXAMPLES = [
    ["examples/teapot.glb", "blue and white porcelain, teapot"],
    ["examples/teapot.glb", "rusty weathered cast iron, teapot"],
    ["examples/suzanne.glb", "monkey head, Sci-Fi digital painting"],
    ["examples/suzanne.glb", "golden fur, cute monkey face, brown eyes"],
]

with gr.Blocks(title="Paint3D UV-Position ControlNet") as demo:
    gr.Markdown(
        "# Paint3D · UV-Position ControlNet\n"
        "Text → lighting-less UV texture for your mesh, using "
        "[GeorgeQi/Paint3d_UVPos_Control](https://huggingface.co/GeorgeQi/Paint3d_UVPos_Control) "
        "(the UV-only pipeline from [Paint3D](https://github.com/OpenTexture/Paint3D), CVPR 2024) on Stable Diffusion 1.5.\n\n"
        "Upload a mesh with UVs (GLB/OBJ). The app rasterizes its **UV position map** and the ControlNet paints the texture "
        "directly in UV space. Works best with clean, large UV islands; heavily fragmented atlases give patchier results."
    )
    with gr.Row():
        with gr.Column(scale=1):
            mesh_in = gr.Model3D(label="Input mesh (.glb / .obj)", clear_color=[0.9, 0.9, 0.9, 1.0])
            prompt = gr.Textbox(label="Prompt", placeholder="e.g. blue and white porcelain, teapot", lines=2)
            run = gr.Button("Paint texture", variant="primary")
            with gr.Accordion("Advanced", open=False):
                neg = gr.Textbox(label="Negative prompt", value=DEFAULT_NEG, lines=2)
                steps = gr.Slider(10, 50, value=30, step=1, label="Steps")
                guidance = gr.Slider(1.0, 12.0, value=7.0, step=0.5, label="Guidance scale")
                cscale = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="ControlNet scale")
                keep_n = gr.Checkbox(value=True, label="Keep input normal map (GLB)")
                seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                rand = gr.Checkbox(value=True, label="Randomize seed")
            pos_btn = gr.Button("Preview UV position map only")
        with gr.Column(scale=1):
            mesh_out = gr.Model3D(label="Textured mesh", clear_color=[0.9, 0.9, 0.9, 1.0])
            with gr.Row():
                tex_out = gr.Image(label="Generated UV texture", type="filepath", format="png")
                pos_out = gr.Image(label="UV position map (condition)", format="png")
            seed_out = gr.Number(label="Seed used", precision=0)

    gr.Examples(
        EXAMPLES, inputs=[mesh_in, prompt], outputs=[mesh_out, tex_out, pos_out, seed_out],
        fn=texture_mesh, cache_examples=True, cache_mode="lazy",
    )
    run.click(texture_mesh, [mesh_in, prompt, neg, steps, guidance, cscale, keep_n, seed, rand],
              [mesh_out, tex_out, pos_out, seed_out])
    pos_btn.click(make_uv_position, [mesh_in], [pos_out])

demo.launch(mcp_server=True)