"""Render a GLB's embedded animation with Blender, including a review contact sheet."""

import argparse
import json
from pathlib import Path
import sys
import bpy
from mathutils import Vector

sys.path.insert(0, str(Path(__file__).resolve().parent))
from render_settings import configure_cpu_render

p = argparse.ArgumentParser()
p.add_argument("--input", required=True)
p.add_argument("--output", required=True)
p.add_argument("--preview", action="store_true")
p.add_argument("--camera-motion", choices=["follow", "fixed"], default="follow")
p.add_argument("--frames", help="Comma-separated exact one-based review frame indices")
p.add_argument("--step", type=int, default=2)
p.add_argument("--frame-count", type=int, default=180)
p.add_argument("--size", type=int, default=640)
p.add_argument("--angle", choices=["quarter", "front", "side", "back", "run"], default="quarter")
a = p.parse_args(sys.argv[sys.argv.index("--") + 1 :])
out = Path(a.output).resolve()
out.mkdir(parents=True, exist_ok=True)
bpy.ops.object.select_all(action="SELECT")
bpy.ops.object.delete(use_global=False)
bpy.context.scene.render.fps = 30
bpy.ops.import_scene.gltf(filepath=str(Path(a.input).resolve()))
scene = bpy.context.scene
configure_cpu_render(scene)
scene.cycles.samples = 12 if scene.cycles.use_denoising else 32
scene.render.resolution_x = scene.render.resolution_y = a.size
scene.render.resolution_percentage = 100
scene.render.image_settings.file_format = "PNG"
scene.world.use_nodes = True
scene.world.node_tree.nodes["Background"].inputs["Color"].default_value = (0.18, 0.21, 0.26, 1)
scene.world.node_tree.nodes["Background"].inputs["Strength"].default_value = 0.65
lights = []
for name, xyz, energy, size in [
    ("key", (-2, -3, 4), 220, 3),
    ("fill", (3, -2, 2), 150, 2.5),
    ("rim", (0, 3, 3), 260, 2),
]:
    d = bpy.data.lights.new(name, "AREA")
    d.energy = energy
    d.shape = "DISK"
    d.size = size
    o = bpy.data.objects.new(name, d)
    scene.collection.objects.link(o)
    o.location = xyz
    lights.append(o)
    o.rotation_euler = (Vector((0, 0, 0.55)) - o.location).to_track_quat("-Z", "Y").to_euler()
bpy.ops.mesh.primitive_plane_add(size=200, location=(0, 0, -0.012))
floor = bpy.context.object
mat = bpy.data.materials.new("Studio floor")
mat.diffuse_color = (0.075, 0.095, 0.13, 1)
floor.data.materials.append(mat)
d = bpy.data.cameras.new("Review Camera")
cam = bpy.data.objects.new("Review Camera", d)
scene.collection.objects.link(cam)
cam.location = {
    "quarter": (1.3, -4, 1.5),
    "front": (0, -4, 0.65),
    "side": (-4, -0.3, 0.65),
    "back": (0, 4, 0.65),
    "run": (3, -4, 1.5),
}[a.angle]
target = Vector((0, 0, 0.54))
cam.rotation_euler = (target - cam.location).to_track_quat("-Z", "Y").to_euler()
d.type = "ORTHO"
d.ortho_scale = 1.45
scene.camera = cam
frames = (
    [1 + round(i * (a.frame_count - 1) / 4) for i in range(5)]
    if a.preview
    else list(range(1, a.frame_count + 1, a.step))
)
if a.frames:
    frames = [int(value) for value in a.frames.split(",")]
    if not frames or any(frame < 1 or frame > a.frame_count for frame in frames):
        p.error("--frames must be within 1..frame-count")
scene.frame_set(1)
rigs = [obj for obj in scene.objects if obj.type == "ARMATURE"]
root_bones = [(obj, bone.name) for obj in rigs for bone in obj.data.bones if bone.parent is None]
if a.camera_motion == "follow" and len(root_bones) != 1:
    raise ValueError("Following the animation requires exactly one root bone")

def root_position():
    obj, name = root_bones[0]
    return obj.matrix_world @ obj.pose.bones[name].head

start_root = root_position() if root_bones else Vector((0, 0, 0))
tracking_objects = [cam, *lights, floor]
initial_locations = [obj.location.copy() for obj in tracking_objects]
camera_frames = []
for i, frame in enumerate(frames):
    scene.frame_set(frame)
    if a.camera_motion == "follow":
        displacement = root_position() - start_root
        displacement.z = 0
        for obj, initial in zip(tracking_objects, initial_locations):
            obj.location = initial + displacement
    camera_frames.append({"frame": frame, "camera": list(cam.location), "floor": list(floor.location)})
    scene.render.filepath = str(out / f"{frame:04d}.png")
    bpy.ops.render.render(write_still=True)
(out / "camera.json").write_text(json.dumps({"policy": a.camera_motion, "frames": camera_frames}, indent=2))
print("RENDERED", len(frames), "frames")