"""Bake a smooth ivory material on rear suit panels while retaining colored trim."""

import argparse
from pathlib import Path
import sys
import bpy

p = argparse.ArgumentParser()
p.add_argument("--input", required=True)
p.add_argument("--output", required=True)
a = p.parse_args(sys.argv[sys.argv.index("--") + 1 :])
out = Path(a.output).resolve()
out.mkdir(parents=True, exist_ok=True)
bpy.ops.wm.open_mainfile(filepath=str(Path(a.input).resolve()))
scene = bpy.context.scene
scene.frame_set(1)
objects = [o for o in scene.objects if o.type == "MESH"]
rig = next(o for o in scene.objects if o.type == "ARMATURE")
rig.data.pose_position = "REST"
source = objects[0].data.materials[0]
bsdf = next(n for n in source.node_tree.nodes if n.type == "BSDF_PRINCIPLED")
original = bsdf.inputs["Base Color"].links[0].from_node.image
mat = bpy.data.materials.new("Rear ivory cleanup bake")
mat.use_nodes = True
nodes = mat.node_tree.nodes
links = mat.node_tree.links
nodes.clear()


def math(op, *values):
    n = nodes.new("ShaderNodeMath")
    n.operation = op
    for i, v in enumerate(values):
        if isinstance(v, (int, float)):
            n.inputs[i].default_value = v
        else:
            links.new(v, n.inputs[i])
    return n.outputs[0]


tex = nodes.new("ShaderNodeTexImage")
tex.image = original
hsv = nodes.new("ShaderNodeSeparateColor")
hsv.mode = "HSV"
links.new(tex.outputs["Color"], hsv.inputs[0])
geom = nodes.new("ShaderNodeNewGeometry")
normal = nodes.new("ShaderNodeSeparateXYZ")
links.new(geom.outputs["Normal"], normal.inputs[0])
pos = nodes.new("ShaderNodeSeparateXYZ")
links.new(geom.outputs["Position"], pos.inputs[0])
rear = math("MINIMUM", math("MAXIMUM", math("MULTIPLY", math("ADD", normal.outputs["Y"], 0.05), 3), 0), 1)
ivory = math("MULTIPLY", math("LESS_THAN", hsv.outputs[0], 0.22), math("LESS_THAN", hsv.outputs[1], 0.60))
ivory = math("MULTIPLY", ivory, math("GREATER_THAN", hsv.outputs[2], 0.10))
mask = math("MULTIPLY", math("MULTIPLY", rear, ivory), math("LESS_THAN", pos.outputs["Z"], 0.76))
mix = nodes.new("ShaderNodeMixRGB")
links.new(mask, mix.inputs[0])
links.new(tex.outputs["Color"], mix.inputs[1])
mix.inputs[2].default_value = (0.72, 0.74, 0.53, 1)
emit = nodes.new("ShaderNodeEmission")
links.new(mix.outputs[0], emit.inputs[0])
output = nodes.new("ShaderNodeOutputMaterial")
links.new(emit.outputs[0], output.inputs["Surface"])
baked = bpy.data.images.new("Reference and clean suit atlas", 4096, 4096, alpha=True)
target = nodes.new("ShaderNodeTexImage")
target.image = baked
nodes.active = target
bpy.ops.object.select_all(action="DESELECT")
for obj in objects:
    obj.data.materials.clear()
    obj.data.materials.append(mat)
    obj.select_set(True)
bpy.context.view_layer.objects.active = objects[0]
bpy.ops.object.duplicate()
copies = list(bpy.context.selected_objects)
bpy.context.view_layer.objects.active = copies[0]
bpy.ops.object.join()
proxy = bpy.context.object
scene.render.engine = "CYCLES"
scene.cycles.samples = 1
scene.render.bake.margin = 16
bpy.ops.object.bake(type="EMIT")
baked.filepath_raw = str(out / "reference-atlas.png")
baked.file_format = "PNG"
baked.save()
bpy.data.objects.remove(proxy, do_unlink=True)
final = bpy.data.materials.new("Reference matched painted suit")
final.use_nodes = True
bsdf = final.node_tree.nodes.get("Principled BSDF")
bsdf.inputs["Roughness"].default_value = 0.75
bsdf.inputs["Metallic"].default_value = 0
tex = final.node_tree.nodes.new("ShaderNodeTexImage")
tex.image = baked
final.node_tree.links.new(tex.outputs["Color"], bsdf.inputs["Base Color"])
for obj in objects:
    obj.data.materials.clear()
    obj.data.materials.append(final)
bpy.ops.object.select_all(action="DESELECT")
for obj in objects:
    obj.select_set(True)
bpy.ops.export_scene.gltf(
    filepath=str(out / "character-static.glb"),
    export_format="GLB",
    use_selection=True,
    export_animations=False,
    export_skins=False,
)
rig.data.pose_position = "POSE"
rig.select_set(True)
bpy.context.view_layer.objects.active = rig
bpy.ops.export_scene.gltf(
    filepath=str(out / "astronaut-wave.glb"),
    export_format="GLB",
    use_selection=True,
    export_animations=True,
    export_skins=True,
    export_frame_range=True,
)
bpy.ops.wm.save_as_mainfile(filepath=str(out / "astronaut-wave.blend"))