"""
3D to Video
────────────────────────────────────────────────────────
Upload a GLB model and turn it into
① a transformed GLB ② an animated GIF preview ③ a metadata JSON.

• Headless-server friendly (EGL + pyglet-headless → pyrender fallback)
• Object in the GIF is now x3 larger (global scale ×3)
• English-only UI with a pastel background
"""

# ──────────────────── 1. Common imports ────────────────────
import os, io, time, glob, json, math, shutil
import numpy as np
from PIL import Image

import pyglet
pyglet.options["headless"] = True
os.environ["PYOPENGL_PLATFORM"] = "egl"

import trimesh
import trimesh.transformations as tf

import gradio as gr
import spaces

LOG_PATH = "./results/demo"
os.makedirs(LOG_PATH, exist_ok=True)

# ──────────────────── 2. Rendering helpers ────────────────────
def _render_with_trimesh(scene: trimesh.Scene, res):
    png = scene.save_image(resolution=res, visible=True)
    if png is None:
        raise RuntimeError("trimesh.save_image returned None")
    return Image.open(io.BytesIO(png)).convert("RGB")

def _render_with_pyrender(mesh_or_scene, res):
    import pyrender
    if isinstance(mesh_or_scene, trimesh.Scene):
        mesh = trimesh.util.concatenate(mesh_or_scene.dump())
    else:
        mesh = mesh_or_scene
    mesh = pyrender.Mesh.from_trimesh(mesh, smooth=False)
    scn = pyrender.Scene()
    scn.add(mesh)
    cam = pyrender.PerspectiveCamera(yfov=np.pi / 3)
    scn.add(cam, pose=tf.translation_matrix([0, 0, 3]))
    light = pyrender.DirectionalLight(intensity=3.0)
    scn.add(light, pose=tf.translation_matrix([0, 5, 5]))
    r = pyrender.OffscreenRenderer(*res)
    color, _ = r.render(scn, flags=pyrender.RenderFlags.RGBA)
    r.delete()
    return Image.fromarray(color[..., :3])

# ──────────────────── 3. GIF generator (object ×3 larger) ────────────────────
def create_model_animation_gif(
    output_path: str,
    input_glb_path: str,
    animation_type: str,
    duration: float = 3.0,
    fps: int = 30,
    resolution=(640, 480),
):
    base = trimesh.load(input_glb_path)
    if isinstance(base, trimesh.Trimesh):
        base = trimesh.Scene(base)

    num_frames = min(int(duration * fps), 60)
    frames = []
    for i in range(num_frames):
        t = i / (num_frames - 1)
        scene = base.copy()

        # -------- enlarge the whole scene by x3 --------
        scene.apply_transform(tf.scale_matrix(3.0))

        # -------- per-frame transform --------
        if animation_type == "rotate":
            M = tf.rotation_matrix(2 * math.pi * t, [0, 1, 0])
        elif animation_type == "float":
            M = tf.translation_matrix([0, 1.5 * math.sin(2 * math.pi * t), 0])  # 0.5*3
        elif animation_type == "pulse":
            M = tf.scale_matrix(0.8 + 0.4 * math.sin(2 * math.pi * t))
        elif animation_type == "explode":
            M = tf.translation_matrix([1.5 * t, 0, 0])  # 0.5*3
        elif animation_type == "assemble":
            M = tf.translation_matrix([1.5 * (1 - t), 0, 0])
        elif animation_type == "swing":
            M = tf.rotation_matrix(math.pi / 6 * math.sin(2 * math.pi * t), [0, 0, 1])
        else:
            M = np.eye(4)
        scene.apply_transform(M)

        # -------- render --------
        try:
            frame = _render_with_trimesh(scene, resolution)
        except Exception as e:
            print("trimesh render failed, pyrender fallback:", e)
            frame = _render_with_pyrender(scene, resolution)

        frames.append(frame)

    frames[0].save(
        output_path,
        save_all=True,
        append_images=frames[1:],
        duration=int(1000 / fps),
        loop=0,
    )
    return output_path

# ──────────────────── 4. Simple GLB transform (unchanged) ────────────────────
def modify_glb_file(input_glb_path, output_glb_path, animation_type="rotate"):
    try:
        scn = trimesh.load(input_glb_path)
        if not isinstance(scn, trimesh.Scene):
            scn = trimesh.Scene(scn)

        if animation_type == "rotate":
            T = tf.rotation_matrix(math.pi / 4, [0, 1, 0])
        elif animation_type == "float":
            T = tf.translation_matrix([0, 0.5, 0])
        elif animation_type == "pulse":
            T = tf.scale_matrix(1.2)
        elif animation_type == "explode":
            T = tf.translation_matrix([0.5, 0, 0])
        elif animation_type == "assemble":
            T = tf.translation_matrix([-0.5, 0, 0])
        elif animation_type == "swing":
            T = tf.rotation_matrix(math.pi / 8, [0, 0, 1])
        else:
            T = np.eye(4)

        scn.apply_transform(T)
        scn.export(output_glb_path)
        return output_glb_path
    except Exception as e:
        print("GLB transform failed, copying original:", e)
        shutil.copy(input_glb_path, output_glb_path)
        return output_glb_path

# ──────────────────── 5. Gradio pipeline ────────────────────
@spaces.GPU
def process_3d_model(input_3d, animation_type, animation_duration, fps):
    try:
        base = os.path.splitext(os.path.basename(input_3d))[0]
        glb_out  = os.path.join(LOG_PATH, f"animated_{base}.glb")
        gif_out  = os.path.join(LOG_PATH, f"preview_{base}.gif")
        json_out = os.path.join(LOG_PATH, f"metadata_{base}.json")

        modify_glb_file(input_3d, glb_out, animation_type)
        create_model_animation_gif(gif_out, input_3d, animation_type,
                                   animation_duration, fps)

        meta = dict(
            animation_type=animation_type,
            duration=animation_duration,
            fps=fps,
            original_model=os.path.basename(input_3d),
            created_at=time.strftime("%Y-%m-%d %H:%M:%S"),
        )
        with open(json_out, "w") as f:
            json.dump(meta, f, indent=4)

        return glb_out, gif_out, json_out

    except Exception as e:
        print("process_3d_model failed:", e)
        err_gif = os.path.join(LOG_PATH, "error.gif")
        Image.new("RGB", (640, 480), (255, 0, 0)).save(err_gif)
        return input_3d, err_gif, None

# ──────────────────── 6. Gradio UI  ────────────────────
PASTEL_CSS = """
body              {background:#f9f6ff !important;}
.gradio-container {background:#f9f6ff !important;}
footer            {display:none !important;}
"""

with gr.Blocks(title="3D to Video", css=PASTEL_CSS) as demo:
    gr.Markdown(
        """
<h2><b>3D to Video</b></h2>
Transform a static GLB model into an animated clip.
Select an animation style, tune duration & FPS, then download the results.
"""
    )
    with gr.Row():
        with gr.Column():
            inp = gr.Model3D(label="Upload GLB")
            typ = gr.Dropdown(
                label="Animation Type",
                choices=["rotate", "float", "explode", "assemble", "pulse", "swing"],
                value="rotate",
            )
            dur = gr.Slider(
                label="Length (seconds)",
                minimum=1.0,
                maximum=10.0,
                value=3.0,
                step=0.5,
            )
            fps = gr.Slider(
                label="FPS",
                minimum=15,
                maximum=60,
                value=30,
                step=1,
            )
            btn = gr.Button("Generate Animation")
        with gr.Column():
            out_glb  = gr.Model3D(label="Animated GLB")
            out_gif  = gr.Image(  label="GIF Preview")
            out_json = gr.File(   label="Metadata JSON")

    btn.click(
        fn=process_3d_model,
        inputs=[inp, typ, dur, fps],
        outputs=[out_glb, out_gif, out_json],
    )

    examples = [[f] for f in glob.glob("./data/demo_glb/*.glb")]
    if examples:
        gr.Examples(examples=examples, inputs=[inp])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)