"""Render seven validation views of a generated 3MF with Blender.

Usage: blender --background --python script/render_previews.py -- model.3mf output_dir
"""

import math
import os
import sys
import zipfile
import xml.etree.ElementTree as ET

import bpy
from mathutils import Vector


def args():
    marker = sys.argv.index("--")
    return os.path.abspath(sys.argv[marker + 1]), os.path.abspath(sys.argv[marker + 2])


source, output = args()
os.makedirs(output, exist_ok=True)
bpy.ops.object.select_all(action="SELECT")
bpy.ops.object.delete(use_global=False)


def import_3mf(path):
    with zipfile.ZipFile(path) as package:
        root = ET.fromstring(package.read("3D/3dmodel.model"))
    namespace = {"m": "http://schemas.microsoft.com/3dmanufacturing/core/2015/02"}
    materials = []
    bases = root.find("m:resources/m:basematerials", namespace)
    if bases is not None:
        for base in bases:
            color = base.attrib.get("displaycolor", "#B0B0B0").lstrip("#")[:6]
            rgba = tuple(int(color[i:i+2], 16) / 255 for i in (0, 2, 4)) + (1.0,)
            material = bpy.data.materials.new(base.attrib.get("name", "material"))
            material.diffuse_color = rgba
            materials.append(material)
    for element in root.findall("m:resources/m:object", namespace):
        mesh_node = element.find("m:mesh", namespace)
        if mesh_node is None:
            continue
        vertices = [(float(v.attrib["x"]), float(v.attrib["y"]), float(v.attrib["z"]))
                    for v in mesh_node.findall("m:vertices/m:vertex", namespace)]
        faces = [(int(t.attrib["v1"]), int(t.attrib["v2"]), int(t.attrib["v3"]))
                 for t in mesh_node.findall("m:triangles/m:triangle", namespace)]
        data = bpy.data.meshes.new(element.attrib.get("name", "mesh"))
        data.from_pydata(vertices, [], faces)
        data.update()
        obj = bpy.data.objects.new(element.attrib.get("name", "object"), data)
        bpy.context.collection.objects.link(obj)
        if materials and "pindex" in element.attrib:
            obj.data.materials.append(materials[int(element.attrib["pindex"])])


import_3mf(source)

objects = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"]
if not objects:
    raise RuntimeError(f"No mesh objects imported from {source}")

corners = []
for obj in objects:
    corners.extend(obj.matrix_world @ Vector(corner) for corner in obj.bound_box)
minimum = Vector((min(p.x for p in corners), min(p.y for p in corners), min(p.z for p in corners)))
maximum = Vector((max(p.x for p in corners), max(p.y for p in corners), max(p.z for p in corners)))
center = (minimum + maximum) / 2
extent = max(maximum - minimum)

world = bpy.context.scene.world
world.color = (0.035, 0.035, 0.035)
bpy.context.scene.render.engine = "BLENDER_EEVEE"
bpy.context.scene.render.resolution_x = 700
bpy.context.scene.render.resolution_y = 700
bpy.context.scene.render.resolution_percentage = 100
bpy.context.scene.render.image_settings.file_format = "PNG"
bpy.context.scene.render.film_transparent = False

camera_data = bpy.data.cameras.new("Validation camera")
camera = bpy.data.objects.new("Validation camera", camera_data)
bpy.context.collection.objects.link(camera)
bpy.context.scene.camera = camera
camera.data.type = "ORTHO"
camera.data.ortho_scale = extent * 1.25

for energy, direction in [(1800, Vector((1, -1, 2))), (900, Vector((-2, 1, 0.5)) )]:
    light_data = bpy.data.lights.new(name="Area", type="AREA")
    light_data.energy = energy
    light_data.shape = "DISK"
    light_data.size = extent
    light = bpy.data.objects.new(name="Area", object_data=light_data)
    light.location = center + direction.normalized() * extent * 2
    bpy.context.collection.objects.link(light)

views = {
    "front": Vector((0, -1, 0)),
    "back": Vector((0, 1, 0)),
    "left": Vector((-1, 0, 0)),
    "right": Vector((1, 0, 0)),
    "top": Vector((0, 0, 1)),
    "bottom": Vector((0, 0, -1)),
    "iso": Vector((1, -1, 0.8)).normalized(),
}
for name, direction in views.items():
    camera.location = center + direction * extent * 3
    camera.rotation_euler = (center - camera.location).to_track_quat("-Z", "Y").to_euler()
    bpy.context.scene.render.filepath = os.path.join(output, f"{name}.png")
    bpy.ops.render.render(write_still=True)

print(f"Rendered {len(views)} views to {output}")