"""Render aligned lid-base and accent STLs with distinct materials."""

import os
import sys
import math

import bpy
from mathutils import Matrix, Vector


def import_mesh(path, name, color):
    bpy.ops.wm.stl_import(filepath=os.path.abspath(path))
    model = bpy.context.selected_objects[0]
    model.name = name
    material = bpy.data.materials.new(f"{name} material")
    material.diffuse_color = color
    material.metallic = 0.0
    material.roughness = 0.45
    model.data.materials.append(material)
    return model


def look_at(camera, target):
    direction = Vector(target) - camera.location
    camera.rotation_euler = direction.to_track_quat("-Z", "Y").to_euler()


def main():
    args = sys.argv[sys.argv.index("--") + 1 :]
    if len(args) not in (4, 5):
        raise SystemExit(
            "expected: lid-base.stl lid-accent.stl output-directory basename [orientation]"
        )
    base_path, accent_path, output_dir, basename = args[:4]
    orientation = args[4] if len(args) == 5 else "usb_left"
    angles = {"usb_left": 0, "usb_front": 90, "usb_right": 180, "usb_rear": 270}
    if orientation not in angles:
        raise SystemExit(f"unknown lid orientation: {orientation}")
    os.makedirs(output_dir, exist_ok=True)
    bpy.ops.wm.read_factory_settings(use_empty=True)
    models = [
        import_mesh(base_path, "Lid base", (0.16, 0.42, 0.72, 1.0)),
        import_mesh(accent_path, "Accent inlays", (0.95, 0.42, 0.08, 1.0)),
    ]
    # Lid exports are print-oriented with the decorated exterior face down.
    # Turn both aligned parts over and present the configured USB edge in its
    # named position so previews inspect the finished face as the user will.
    for model in models:
        model.matrix_world = (
            Matrix.Rotation(math.radians(angles[orientation]), 4, "Z")
            @ Matrix.Rotation(math.pi, 4, "X")
        )
    bpy.context.view_layer.update()

    corners = [obj.matrix_world @ Vector(corner) for obj in models for corner in obj.bound_box]
    mins = Vector([min(v[i] for v in corners) for i in range(3)])
    maxs = Vector([max(v[i] for v in corners) for i in range(3)])
    center = (mins + maxs) / 2
    span = max(maxs - mins)

    world = bpy.data.worlds.new("Inspection world")
    world.use_nodes = True
    world.node_tree.nodes["Background"].inputs["Color"].default_value = (0.035, 0.04, 0.05, 1)
    world.node_tree.nodes["Background"].inputs["Strength"].default_value = 0.5
    bpy.context.scene.world = world

    camera_data = bpy.data.cameras.new("Camera")
    camera = bpy.data.objects.new("Camera", camera_data)
    bpy.context.collection.objects.link(camera)
    bpy.context.scene.camera = camera
    camera.data.type = "ORTHO"
    camera.data.ortho_scale = span * 1.35

    key_data = bpy.data.lights.new("Key", type="AREA")
    key_data.energy = 900
    key_data.size = span
    key = bpy.data.objects.new("Key", key_data)
    key.location = center + Vector((span, -span, span * 1.5))
    bpy.context.collection.objects.link(key)

    scene = bpy.context.scene
    scene.render.engine = "BLENDER_WORKBENCH"
    scene.render.resolution_x = 800
    scene.render.resolution_y = 600
    scene.render.resolution_percentage = 100
    scene.render.image_settings.file_format = "PNG"
    scene.display.shading.light = "STUDIO"
    scene.display.shading.color_type = "MATERIAL"
    scene.display.shading.show_shadows = True
    scene.display.shading.show_cavity = True
    scene.display.shading.cavity_type = "BOTH"

    for view, direction in {"iso": (1, -1, 0.8), "top": (0, 0, 1)}.items():
        camera.location = center + Vector(direction).normalized() * span * 2.2
        look_at(camera, center)
        scene.render.filepath = os.path.join(output_dir, f"{basename}_{view}.png")
        bpy.ops.render.render(write_still=True)


if __name__ == "__main__":
    main()