#!/usr/bin/env python3
"""Build and render web-ready polished-silver GLB variants with Blender."""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import bpy
import bmesh
from mathutils import Vector


VARIANTS = {
    "source": {"target_triangles": None, "replace_material": False, "export": False},
    "review": {"target_triangles": None, "replace_material": True, "export": False},
    "web": {"target_triangles": 280_000, "replace_material": True, "export": True},
    "mobile": {"target_triangles": 90_000, "replace_material": True, "export": True},
}


def parse_args() -> argparse.Namespace:
    argv = []
    if "--" in __import__("sys").argv:
        argv = __import__("sys").argv[__import__("sys").argv.index("--") + 1 :]
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True)
    parser.add_argument("--outdir", required=True)
    parser.add_argument("--variant", required=True, choices=VARIANTS)
    return parser.parse_args(argv)


def clear_scene() -> None:
    bpy.ops.object.select_all(action="SELECT")
    bpy.ops.object.delete(use_global=False)
    for datablocks in (
        bpy.data.meshes,
        bpy.data.curves,
        bpy.data.materials,
        bpy.data.cameras,
        bpy.data.lights,
    ):
        for datablock in list(datablocks):
            if datablock.users == 0:
                datablocks.remove(datablock)


def import_and_prepare(path: Path) -> bpy.types.Object:
    bpy.ops.import_scene.gltf(filepath=str(path))
    meshes = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"]
    if not meshes:
        raise RuntimeError("No mesh object was imported")

    bpy.ops.object.select_all(action="DESELECT")
    for obj in meshes:
        obj.select_set(True)
    bpy.context.view_layer.objects.active = meshes[0]
    if len(meshes) > 1:
        bpy.ops.object.join()
    obj = bpy.context.view_layer.objects.active

    bpy.ops.object.transform_apply(location=False, rotation=True, scale=True)
    obj.data.validate(verbose=False, clean_customdata=False)
    for polygon in obj.data.polygons:
        polygon.use_smooth = True

    corners = [obj.matrix_world @ Vector(corner) for corner in obj.bound_box]
    center = sum(corners, Vector()) / 8.0
    obj.location -= center
    bpy.context.view_layer.update()
    obj.name = "TN_Silver_Product"
    obj.data.name = "TN_Silver_Product_Mesh"
    return obj


def triangle_count(obj: bpy.types.Object) -> int:
    obj.data.calc_loop_triangles()
    return len(obj.data.loop_triangles)


def decimate_to(obj: bpy.types.Object, target: int | None) -> None:
    if not target:
        return
    current = triangle_count(obj)
    if current <= target:
        return
    modifier = obj.modifiers.new(name="Web silhouette optimization", type="DECIMATE")
    modifier.decimate_type = "COLLAPSE"
    modifier.ratio = max(0.001, min(1.0, target / current))
    modifier.use_collapse_triangulate = True
    bpy.context.view_layer.objects.active = obj
    bpy.ops.object.modifier_apply(modifier=modifier.name)
    for polygon in obj.data.polygons:
        polygon.use_smooth = True


def smooth_generated_surface(obj: bpy.types.Object) -> None:
    modifier = obj.modifiers.new(name="Generated surface cleanup", type="LAPLACIANSMOOTH")
    modifier.lambda_factor = 0.12
    modifier.lambda_border = 0.02
    modifier.iterations = 5
    modifier.use_volume_preserve = True
    modifier.use_normalized = True
    bpy.context.view_layer.objects.active = obj
    bpy.ops.object.modifier_apply(modifier=modifier.name)


def weld_generated_seams(obj: bpy.types.Object) -> None:
    dimensions = obj.dimensions
    weld_distance = max(dimensions.x, dimensions.y, dimensions.z) * 1e-6
    mesh = obj.data
    editable = bmesh.new()
    editable.from_mesh(mesh)
    bmesh.ops.remove_doubles(editable, verts=list(editable.verts), dist=weld_distance)
    bmesh.ops.recalc_face_normals(editable, faces=list(editable.faces))
    editable.to_mesh(mesh)
    editable.free()
    mesh.update()
    for layer in list(mesh.uv_layers):
        mesh.uv_layers.remove(layer)


def polished_silver_material() -> bpy.types.Material:
    material = bpy.data.materials.new("TN Polished Silver")
    material.use_nodes = True
    material.use_backface_culling = True
    material.diffuse_color = (0.972, 0.960, 0.915, 1.0)
    node = material.node_tree.nodes.get("Principled BSDF")
    node.inputs["Base Color"].default_value = (0.972, 0.960, 0.915, 1.0)
    node.inputs["Metallic"].default_value = 1.0
    node.inputs["Roughness"].default_value = 0.055
    if "Coat Weight" in node.inputs:
        node.inputs["Coat Weight"].default_value = 0.0
    return material


def replace_material(obj: bpy.types.Object) -> None:
    material = polished_silver_material()
    obj.data.materials.clear()
    obj.data.materials.append(material)


def object_bounds(obj: bpy.types.Object) -> tuple[Vector, Vector, Vector, float]:
    corners = [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) * 0.5
    radius = max((maximum - minimum).length * 0.5, 0.1)
    return minimum, maximum, center, radius


def point_at(obj: bpy.types.Object, target: Vector) -> None:
    obj.rotation_euler = (target - obj.location).to_track_quat("-Z", "Y").to_euler()


def add_area_light(
    name: str,
    location: tuple[float, float, float],
    energy: float,
    size: float,
    color: tuple[float, float, float],
    target: Vector,
) -> None:
    data = bpy.data.lights.new(name=name, type="AREA")
    data.energy = energy
    data.shape = "DISK"
    data.size = size
    data.color = color
    light = bpy.data.objects.new(name, data)
    bpy.context.collection.objects.link(light)
    light.location = location
    point_at(light, target)


def setup_studio(obj: bpy.types.Object, output_path: Path) -> None:
    scene = bpy.context.scene
    scene.render.engine = "BLENDER_EEVEE"
    scene.render.resolution_x = 1000
    scene.render.resolution_y = 1000
    scene.render.resolution_percentage = 100
    scene.render.image_settings.file_format = "PNG"
    scene.render.film_transparent = False
    scene.render.filepath = str(output_path)
    scene.render.image_settings.color_mode = "RGBA"
    scene.render.use_file_extension = True
    scene.render.resolution_percentage = 100
    scene.render.image_settings.color_depth = "8"

    scene.world.use_nodes = True
    world_nodes = scene.world.node_tree.nodes
    world_links = scene.world.node_tree.links
    world_nodes.clear()
    environment = world_nodes.new("ShaderNodeTexEnvironment")
    environment.image = bpy.data.images.load(
        str(Path(__file__).with_name("highkey-studio-env.png")), check_existing=True
    )
    environment.interpolation = "Linear"
    background = world_nodes.new("ShaderNodeBackground")
    background.inputs["Strength"].default_value = 0.82
    world_output = world_nodes.new("ShaderNodeOutputWorld")
    world_links.new(environment.outputs["Color"], background.inputs["Color"])
    world_links.new(background.outputs["Background"], world_output.inputs["Surface"])

    try:
        scene.view_settings.look = "AgX - Medium High Contrast"
    except TypeError:
        pass
    scene.view_settings.exposure = 0.55

    minimum, maximum, center, radius = object_bounds(obj)
    camera_data = bpy.data.cameras.new("Studio Camera")
    camera = bpy.data.objects.new("Studio Camera", camera_data)
    bpy.context.collection.objects.link(camera)
    camera.location = center + Vector((0.0, -2.45 * radius, 2.65 * radius))
    camera.data.lens = 58
    camera.data.sensor_width = 36
    point_at(camera, center + Vector((0.0, 0.0, -0.03 * radius)))
    scene.camera = camera

    add_area_light(
        "Key softbox",
        tuple(center + Vector((2.8 * radius, -2.7 * radius, 3.8 * radius))),
        1150,
        2.8 * radius,
        (1.0, 0.94, 0.88),
        center,
    )
    add_area_light(
        "Fill softbox",
        tuple(center + Vector((-3.4 * radius, -0.8 * radius, 2.1 * radius))),
        850,
        3.7 * radius,
        (0.76, 0.86, 1.0),
        center,
    )
    add_area_light(
        "Edge strip",
        tuple(center + Vector((1.1 * radius, 3.0 * radius, 3.0 * radius))),
        1450,
        1.5 * radius,
        (1.0, 1.0, 1.0),
        center,
    )

    plane_size = radius * 8.0
    bpy.ops.mesh.primitive_plane_add(size=plane_size, location=(center.x, center.y, minimum.z - 0.055 * radius))
    floor = bpy.context.object
    floor.name = "Studio Floor"
    floor_material = bpy.data.materials.new("Studio Floor Material")
    floor_material.use_nodes = True
    floor_bsdf = floor_material.node_tree.nodes.get("Principled BSDF")
    floor_bsdf.inputs["Base Color"].default_value = (0.42, 0.44, 0.48, 1.0)
    floor_bsdf.inputs["Roughness"].default_value = 0.27
    floor.data.materials.append(floor_material)


def export_glb(obj: bpy.types.Object, output_path: Path) -> None:
    bpy.ops.object.select_all(action="DESELECT")
    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj
    bpy.ops.export_scene.gltf(
        filepath=str(output_path),
        export_format="GLB",
        use_selection=True,
        export_apply=True,
        export_yup=True,
        export_normals=True,
        export_tangents=False,
        export_texcoords=False,
        export_materials="EXPORT",
        export_animations=False,
        export_cameras=False,
        export_lights=False,
        export_draco_mesh_compression_enable=False,
    )


def main() -> None:
    args = parse_args()
    source_path = Path(args.input).resolve()
    output_dir = Path(args.outdir).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    config = VARIANTS[args.variant]

    clear_scene()
    obj = import_and_prepare(source_path)
    source_triangles = triangle_count(obj)
    if config["replace_material"]:
        weld_generated_seams(obj)
        smooth_generated_surface(obj)
    decimate_to(obj, config["target_triangles"])
    if config["replace_material"]:
        replace_material(obj)
    final_triangles = triangle_count(obj)

    glb_path = None
    if config["export"]:
        glb_path = output_dir / f"model-silver-{args.variant}.glb"
        export_glb(obj, glb_path)

    preview_path = output_dir / f"preview-{args.variant}.png"
    setup_studio(obj, preview_path)
    bpy.ops.render.render(write_still=True)

    minimum, maximum, _, _ = object_bounds(obj)
    result = {
        "variant": args.variant,
        "source": str(source_path),
        "source_triangles": source_triangles,
        "triangles": final_triangles,
        "vertices": len(obj.data.vertices),
        "dimensions": [round(value, 6) for value in (maximum - minimum)],
        "glb": str(glb_path) if glb_path else None,
        "glb_bytes": glb_path.stat().st_size if glb_path else None,
        "preview": str(preview_path),
        "material": "source" if not config["replace_material"] else {
            "name": "TN Polished Silver",
            "base_color": [0.972, 0.960, 0.915, 1.0],
            "metallic": 1.0,
            "roughness": 0.055,
        },
    }
    manifest_path = output_dir / f"manifest-{args.variant}.json"
    manifest_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
    print("TN_RESULT", json.dumps(result))


if __name__ == "__main__":
    main()