"""Blender scene builder for the optimized four-mirror train.

Run headless, from the directory containing the geometry JSON:

    blender --background --python emberscope-four-mirror-optimized-v1.blender_scene.py -- \
        --geometry emberscope-four-mirror-optimized-v1.blender.json --out-dir renders --resolution 1600

The CI/CD deploy build runs this rendering automatically on raksasa through
scripts/render_blender_site_assets.sh. You only need to commit generated PNGs
when you want to inspect or compare those images in git.

Builds gold mirror meshes, an emissive folded beam, the detector plate, and
the 150 mm envelope wireframe, then renders a hero view plus side, top,
front, and reverse inspection views.
"""

import argparse
import json
import math
import sys
from pathlib import Path

import bpy
from mathutils import Vector


def parse_args():
    argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    parser = argparse.ArgumentParser()
    parser.add_argument("--geometry", default="emberscope-four-mirror-optimized-v1.blender.json")
    parser.add_argument("--out-dir", default="renders")
    parser.add_argument("--resolution", type=int, default=1600)
    return parser.parse_args(argv)


def make_material(name, color, metallic, roughness, emission=0.0):
    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    bsdf = mat.node_tree.nodes["Principled BSDF"]
    if emission > 0.0:
        # Pure emitter: black base so the beam colour stays saturated.
        bsdf.inputs["Base Color"].default_value = (0.0, 0.0, 0.0, 1.0)
        bsdf.inputs["Emission Color"].default_value = (*color, 1.0)
        bsdf.inputs["Emission Strength"].default_value = emission
    else:
        bsdf.inputs["Base Color"].default_value = (*color, 1.0)
    bsdf.inputs["Metallic"].default_value = metallic
    bsdf.inputs["Roughness"].default_value = roughness
    return mat


def add_mesh(name, vertices, faces, material):
    mesh = bpy.data.meshes.new(name)
    mesh.from_pydata([Vector(v) for v in vertices], [], faces)
    mesh.update()
    obj = bpy.data.objects.new(name, mesh)
    obj.data.materials.append(material)
    bpy.context.collection.objects.link(obj)
    return obj


def add_polyline(name, points, material, radius=0.18):
    curve = bpy.data.curves.new(name, type="CURVE")
    curve.dimensions = "3D"
    curve.bevel_depth = radius
    spline = curve.splines.new("POLY")
    spline.points.add(len(points) - 1)
    for sp, p in zip(spline.points, points):
        sp.co = (p[0], p[1], p[2], 1.0)
    obj = bpy.data.objects.new(name, curve)
    obj.data.materials.append(material)
    bpy.context.collection.objects.link(obj)
    return obj


def main():
    args = parse_args()
    geometry = json.loads(Path(args.geometry).read_text())
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    bpy.ops.object.select_all(action="SELECT")
    bpy.ops.object.delete()

    gold = make_material("mirror_gold", (0.89, 0.62, 0.19), 1.0, 0.30)
    dark = make_material("plate_dark", (0.05, 0.06, 0.06), 0.2, 0.5)
    wire = make_material("wire", (0.40, 0.48, 0.45), 0.0, 0.8, emission=0.5)
    beam_colors = {
        -1.5: (0.85, 0.28, 0.07),
        0.0: (0.95, 0.60, 0.10),
        1.5: (0.04, 0.52, 0.44),
    }

    centers = []
    for mirror in geometry["mirrors"]:
        obj = add_mesh(mirror["name"], mirror["vertices"], mirror["faces"], gold)
        mod = obj.modifiers.new("thickness", "SOLIDIFY")
        mod.thickness = geometry.get("mirror_thickness_mm", 4.0)
        centers.extend(mirror["vertices"][::40])

    quad = geometry["sensor_quad_mm"]
    add_mesh("detector", quad, [[0, 1, 2, 3]], dark)

    # Entrance aperture ring so the truncated entry rays read as a stop.
    ap = geometry["aperture"]
    ring_pts = []
    radius = ap["diameter_mm"] / 2.0 + 1.0
    for i in range(49):
        theta = 2.0 * math.pi * i / 48
        ring_pts.append(
            (
                ap["center_mm"][0] + radius * math.cos(theta),
                ap["center_mm"][1] + radius * math.sin(theta),
                ap["center_mm"][2],
            )
        )
    add_polyline("aperture_ring", ring_pts, wire, radius=0.5)

    for fdata in geometry["beam_fields"]:
        color = beam_colors.get(fdata["field_deg"], (0.9, 0.7, 0.3))
        beam_mat = make_material(
            f"beam_{fdata['field_deg']}", color, 0.0, 0.4, emission=7.0
        )
        for i, ray in enumerate(fdata["rays"]):
            add_polyline(f"ray_{fdata['field_deg']}_{i}", ray, beam_mat, radius=0.11)

    cc = Vector(geometry["envelope_cube_center_mm"])
    half = geometry["envelope_target_mm"] / 2.0
    corners = []
    for sx in (-1, 1):
        for sy in (-1, 1):
            for sz in (-1, 1):
                corners.append(cc + Vector((sx * half, sy * half, sz * half)))
    edges = [
        (0, 1), (2, 3), (4, 5), (6, 7),
        (0, 2), (1, 3), (4, 6), (5, 7),
        (0, 4), (1, 5), (2, 6), (3, 7),
    ]
    for index, (a, b) in enumerate(edges):
        add_polyline(f"cube_edge_{index}", [corners[a], corners[b]], wire, radius=0.22)

    bounds_min = Vector(geometry["geometry_bounds_mm"]["min"])
    bounds_max = Vector(geometry["geometry_bounds_mm"]["max"])
    center = (bounds_min + bounds_max) / 2.0
    span = max(bounds_max - bounds_min)

    world = bpy.context.scene.world or bpy.data.worlds.new("world")
    bpy.context.scene.world = world
    world.use_nodes = True
    world.node_tree.nodes["Background"].inputs["Color"].default_value = (
        0.020, 0.027, 0.025, 1.0,
    )

    def add_light(name, energy, offset):
        light = bpy.data.objects.new(name, bpy.data.lights.new(name, "AREA"))
        light.data.energy = energy
        light.data.size = span * 1.5
        light.location = center + Vector(offset)
        direction = center - light.location
        light.rotation_euler = direction.to_track_quat("-Z", "Y").to_euler()
        bpy.context.collection.objects.link(light)

    add_light("key", 1.4e6, (span * 1.3, -span * 0.9, span * 1.2))
    add_light("fill", 5.0e5, (-span * 1.5, span * 0.5, span * 0.7))
    add_light("rim", 3.0e5, (-span * 0.3, -span * 1.4, -span * 0.9))

    cam_data = bpy.data.cameras.new("camera")
    camera = bpy.data.objects.new("camera", cam_data)
    bpy.context.collection.objects.link(camera)
    bpy.context.scene.camera = camera

    def aim(location, ortho=None):
        camera.location = location
        direction = center - Vector(location)
        camera.rotation_euler = direction.to_track_quat("-Z", "Y").to_euler()
        if ortho:
            cam_data.type = "ORTHO"
            cam_data.ortho_scale = ortho
        else:
            cam_data.type = "PERSP"
            cam_data.lens = 55.0

    scene = bpy.context.scene
    engines = {e.identifier for e in
               type(scene.render).bl_rna.properties["engine"].enum_items}
    for engine in ("BLENDER_EEVEE_NEXT", "BLENDER_EEVEE", "CYCLES"):
        if engine in engines:
            scene.render.engine = engine
            break
    if scene.render.engine == "CYCLES":
        scene.cycles.samples = 64
    try:
        # AgX desaturates the beam colours badly; Standard keeps them.
        scene.view_settings.view_transform = "Standard"
    except TypeError:
        pass
    scene.render.resolution_x = args.resolution
    scene.render.resolution_y = int(args.resolution * 0.625)
    scene.render.film_transparent = False

    views = {
        "hero": (center + Vector((span * 1.8, -span * 0.9, span * 0.6)), None),
        "side": (center + Vector((span * 2.2, 0.0, 0.0)), span * 1.35),
        "top": (center + Vector((0.0, 0.0, span * 2.2)), span * 1.35),
        "front": (center + Vector((0.0, -span * 2.2, 0.0)), span * 1.35),
        "reverse": (center + Vector((-span * 1.8, span * 0.9, span * 0.6)), None),
    }
    for name, (location, ortho) in views.items():
        aim(location, ortho)
        scene.render.filepath = str(out_dir / f"four-mirror-train-{name}.png")
        bpy.ops.render.render(write_still=True)


main()