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

import argparse
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("--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
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
    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))
)
for i, frame in enumerate(frames):
    scene.frame_set(frame)
    scene.render.filepath = str(out / f"{frame:04d}.png")
    bpy.ops.render.render(write_still=True)
print("RENDERED", len(frames), "frames")