"""Project the 04:27 reference onto the wearer and bake it into the GLB UVs."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import bpy


def math_node(nodes, operation: str, value: float):
    node = nodes.new("ShaderNodeMath")
    node.operation = operation
    node.inputs[1].default_value = value
    return node


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", required=True, type=Path)
    parser.add_argument("--reference", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--size", type=int, default=2048)
    argv = sys.argv[sys.argv.index("--") + 1 :] if "--" in sys.argv else []
    args = parser.parse_args(argv)

    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.ops.import_scene.gltf(filepath=str(args.model))
    meshes = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"]
    if len(meshes) != 1:
        raise RuntimeError(f"Expected one mesh, found {len(meshes)}")
    wearer = meshes[0]

    bpy.ops.object.select_all(action="DESELECT")
    wearer.select_set(True)
    bpy.context.view_layer.objects.active = wearer
    bpy.ops.object.transform_apply(location=False, rotation=True, scale=True)

    if not wearer.data.uv_layers:
        raise RuntimeError("The wearer GLB has no UV map")
    wearer.data.uv_layers.active_index = 0

    source = bpy.data.images.load(str(args.reference), check_existing=False)
    source.colorspace_settings.name = "sRGB"
    baked = bpy.data.images.new(
        "WearerColorBaked",
        width=args.size,
        height=args.size,
        alpha=True,
        float_buffer=False,
    )
    baked.colorspace_settings.name = "sRGB"
    baked.generated_color = (0.70, 0.63, 0.56, 1.0)

    material = bpy.data.materials.new("Wearer projected color")
    material.use_nodes = True
    nodes = material.node_tree.nodes
    links = material.node_tree.links
    nodes.clear()

    output = nodes.new("ShaderNodeOutputMaterial")
    emission = nodes.new("ShaderNodeEmission")
    coordinates = nodes.new("ShaderNodeTexCoord")
    separate = nodes.new("ShaderNodeSeparateXYZ")
    combine = nodes.new("ShaderNodeCombineXYZ")
    source_texture = nodes.new("ShaderNodeTexImage")
    source_texture.image = source
    source_texture.interpolation = "Linear"
    source_texture.extension = "CLIP"
    target_texture = nodes.new("ShaderNodeTexImage")
    target_texture.image = baked
    target_texture.select = True
    nodes.active = target_texture

    # Generated X maps the model width. Generated Z maps bottom-to-top after
    # applying the imported GLB transform. The narrower ranges align the mesh
    # bounds to the wearer inside the full 04:27 photograph.
    scale_x = math_node(nodes, "MULTIPLY", 0.72)
    offset_x = math_node(nodes, "ADD", 0.14)
    scale_y = math_node(nodes, "MULTIPLY", 0.83)
    offset_y = math_node(nodes, "ADD", 0.14)

    links.new(coordinates.outputs["Generated"], separate.inputs["Vector"])
    links.new(separate.outputs["X"], scale_x.inputs[0])
    links.new(scale_x.outputs[0], offset_x.inputs[0])
    links.new(offset_x.outputs[0], combine.inputs["X"])
    links.new(separate.outputs["Z"], scale_y.inputs[0])
    links.new(scale_y.outputs[0], offset_y.inputs[0])
    links.new(offset_y.outputs[0], combine.inputs["Y"])
    links.new(combine.outputs["Vector"], source_texture.inputs["Vector"])
    links.new(source_texture.outputs["Color"], emission.inputs["Color"])
    links.new(emission.outputs["Emission"], output.inputs["Surface"])

    wearer.data.materials.clear()
    wearer.data.materials.append(material)

    scene = bpy.context.scene
    scene.render.image_settings.file_format = "PNG"
    scene.render.image_settings.color_mode = "RGBA"
    scene.render.image_settings.color_depth = "8"
    scene.render.film_transparent = True

    # Image baking is implemented by Cycles.
    scene.render.engine = "CYCLES"
    scene.cycles.device = "CPU"
    scene.cycles.samples = 1
    bpy.ops.object.bake(type="EMIT", margin=24, use_clear=True)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    baked.filepath_raw = str(args.output)
    baked.file_format = "PNG"
    baked.save()
    print(f"saved={args.output} size={args.size}x{args.size}")


if __name__ == "__main__":
    main()