import bpy
import math
import mathutils
import sys
from pathlib import Path


input_path = Path(sys.argv[sys.argv.index("--") + 1])
output_dir = Path(sys.argv[sys.argv.index("--") + 2])
single_render = len(sys.argv) > sys.argv.index("--") + 3 and sys.argv[sys.argv.index("--") + 3] == "single"
output_dir.mkdir(parents=True, exist_ok=True)

bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.ops.import_scene.gltf(filepath=str(input_path))
scene = bpy.context.scene
ring = next(obj for obj in scene.objects if obj.type == "MESH")
mesh = ring.data
original_corner_normals = [normal.vector.copy() for normal in mesh.corner_normals]

target_size = 2.05
scale = target_size / max(ring.dimensions)
ring.scale = tuple(component * scale for component in ring.scale)
ring.rotation_mode = "XYZ"
ring.rotation_euler = (-1.18, -0.05, math.pi - 0.06)
bpy.context.view_layer.update()

silver = bpy.data.materials.new("Polished sterling silver")
silver.use_nodes = True
principled = silver.node_tree.nodes.get("Principled BSDF")
principled.inputs["Base Color"].default_value = (0.91, 0.925, 0.94, 1.0)
principled.inputs["Metallic"].default_value = 1.0
principled.inputs["Roughness"].default_value = 0.16
ring.data.materials.clear()
ring.data.materials.append(silver)

ground_material = bpy.data.materials.new("Black sweep")
ground_material.use_nodes = True
ground_principled = ground_material.node_tree.nodes.get("Principled BSDF")
ground_principled.inputs["Base Color"].default_value = (0.004, 0.005, 0.007, 1.0)
ground_principled.inputs["Roughness"].default_value = 0.30

minimum_z = min((ring.matrix_world @ mathutils.Vector(corner)).z for corner in ring.bound_box)
bpy.ops.mesh.primitive_plane_add(size=20, location=(0, 0, minimum_z - 0.015))
ground = bpy.context.object
ground.data.materials.append(ground_material)

camera_data = bpy.data.cameras.new("Camera")
camera = bpy.data.objects.new("Camera", camera_data)
scene.collection.objects.link(camera)
scene.camera = camera
camera.location = (0, 0, 5.05)
camera.rotation_euler = (0, 0, 0)
camera.data.angle = math.radians(30)


def point_at(obj, target=(0, 0, 0)):
    direction = mathutils.Vector(target) - obj.location
    obj.rotation_euler = direction.to_track_quat("-Z", "Y").to_euler()


def area_light(name, location, energy, width, height, color):
    light_data = bpy.data.lights.new(name, type="AREA")
    light_data.energy = energy
    light_data.shape = "RECTANGLE"
    light_data.size = width
    light_data.size_y = height
    light_data.color = color
    light = bpy.data.objects.new(name, light_data)
    scene.collection.objects.link(light)
    light.location = location
    point_at(light)
    return light


# Long, narrow sources create the readable reflection bands used in jewelry photography.
area_light("Key strip", (-3.4, -1.8, 3.7), 1150, 4.6, 0.65, (1.0, 0.93, 0.84))
area_light("Right strip", (3.6, 0.6, 2.8), 920, 3.5, 0.38, (0.79, 0.91, 1.0))
area_light("Rear rim", (0.3, 3.5, 1.9), 1250, 4.8, 0.28, (0.86, 0.93, 1.0))
area_light("Front fill", (-0.2, -3.0, 1.35), 420, 3.0, 1.1, (1.0, 0.96, 0.90))

world = bpy.data.worlds.new("Dark studio")
scene.world = world
world.use_nodes = True
background = world.node_tree.nodes.get("Background")
background.inputs["Color"].default_value = (0.002, 0.0025, 0.0035, 1.0)
background.inputs["Strength"].default_value = 0.10

scene.render.engine = "BLENDER_EEVEE"
scene.render.resolution_x = 1200
scene.render.resolution_y = 800
scene.render.resolution_percentage = 100
scene.render.image_settings.file_format = "PNG"
scene.render.image_settings.color_mode = "RGBA"
scene.render.film_transparent = False
scene.render.image_settings.color_depth = "8"
scene.render.resolution_percentage = 100
scene.render.filepath = str(output_dir / "studio-original-normals.png")
try:
    scene.view_settings.look = "AgX - Medium High Contrast"
except TypeError:
    pass

bpy.ops.render.render(write_still=True)
if single_render:
    raise SystemExit(0)

# Match the area-weighted vertex-normal recomputation performed by Three.js.
mesh.calc_loop_triangles()
area_weighted_normals = [mathutils.Vector((0.0, 0.0, 0.0)) for _ in mesh.vertices]
for tri in mesh.loop_triangles:
    weighted = tri.normal * tri.area
    for vertex_index in tri.vertices:
        area_weighted_normals[vertex_index] += weighted
for normal in area_weighted_normals:
    if normal.length_squared:
        normal.normalize()
mesh.normals_split_custom_set_from_vertices(area_weighted_normals)
mesh.update()
scene.render.filepath = str(output_dir / "studio-recomputed-normals.png")
bpy.ops.render.render(write_still=True)

# A matte diagnostic separates true geometric waviness from metallic reflections.
mesh.normals_split_custom_set(original_corner_normals)
mesh.update()
principled.inputs["Base Color"].default_value = (0.24, 0.255, 0.27, 1.0)
principled.inputs["Metallic"].default_value = 0.0
principled.inputs["Roughness"].default_value = 0.42
scene.render.filepath = str(output_dir / "matte-geometry-diagnostic.png")
bpy.ops.render.render(write_still=True)