"""Bake the front reference into the existing UV atlas, retaining unseen source texture."""

import argparse
from pathlib import Path
import sys
import json
import bpy
import numpy as np
from mathutils import Vector
from mathutils.bvhtree import BVHTree

p = argparse.ArgumentParser()
p.add_argument("--mesh", required=True)
p.add_argument("--reference", required=True)
p.add_argument("--output", required=True)
a = p.parse_args(sys.argv[sys.argv.index("--") + 1 :])
out = Path(a.output)
out.mkdir(parents=True, exist_ok=True)
bpy.ops.object.select_all(action="SELECT")
bpy.ops.object.delete(use_global=False)
bpy.ops.import_scene.gltf(filepath=str(Path(a.mesh).resolve()))
objects = [o for o in bpy.context.scene.objects if o.type == "MESH"]
for obj in objects:
    bpy.context.view_layer.objects.active = obj
    obj.select_set(True)
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    obj.select_set(False)
vertices = []
triangles = []
offset = 0
for obj in objects:
    obj.data.calc_loop_triangles()
    vertices.extend(tuple(v.co) for v in obj.data.vertices)
    triangles.extend(tuple(offset + i for i in t.vertices) for t in obj.data.loop_triangles)
    offset += len(obj.data.vertices)
vertices = np.array(vertices)
lo = vertices.min(0)
hi = vertices.max(0)
ref = bpy.data.images.load(str(Path(a.reference).resolve()))
w, h = ref.size
pixels = np.array(ref.pixels[:]).reshape(h, w, 4)
mask = pixels[:, :, :3].min(2) < 0.88
ys, xs = np.where(mask)
uvmin = np.array([xs.min() / w, ys.min() / h])
uvmax = np.array([xs.max() / w, ys.max() / h])
sx = (uvmax[0] - uvmin[0]) / (hi[0] - lo[0])
sz = (uvmax[1] - uvmin[1]) / (hi[2] - lo[2])
tx = uvmin[0] - lo[0] * sx
tz = uvmin[1] - lo[2] * sz
bvh = BVHTree.FromPolygons(vertices.tolist(), triangles, all_triangles=True)
old_mat = objects[0].data.materials[0]
old_tex = (
    next(
        n.image
        for n in old_mat.node_tree.nodes
        if n.type == "TEX_IMAGE" and n.image and "baseColor" in n.image.name
    )
    if any(n.type == "TEX_IMAGE" and n.image and "baseColor" in n.image.name for n in old_mat.node_tree.nodes)
    else next(n.image for n in old_mat.node_tree.nodes if n.type == "TEX_IMAGE" and n.image)
)
# Resolve the actual base-color connection instead of relying on material node order.
bsdf = next(n for n in old_mat.node_tree.nodes if n.type == "BSDF_PRINCIPLED")
if bsdf.inputs["Base Color"].is_linked:
    old_tex = bsdf.inputs["Base Color"].links[0].from_node.image
mat = bpy.data.materials.new("Reference projection bake")
mat.use_nodes = True
nodes = mat.node_tree.nodes
links = mat.node_tree.links
nodes.clear()
output = nodes.new("ShaderNodeOutputMaterial")
emit = nodes.new("ShaderNodeEmission")
links.new(emit.outputs[0], output.inputs["Surface"])
original = nodes.new("ShaderNodeTexImage")
original.image = old_tex
projected = nodes.new("ShaderNodeTexImage")
projected.image = ref
projected.extension = "CLIP"
geom = nodes.new("ShaderNodeNewGeometry")
sep = nodes.new("ShaderNodeSeparateXYZ")
links.new(geom.outputs["Position"], sep.inputs[0])
combine = nodes.new("ShaderNodeCombineXYZ")
for channel, scale, translate in [("X", sx, tx), ("Z", sz, tz)]:
    mult = nodes.new("ShaderNodeMath")
    mult.operation = "MULTIPLY_ADD"
    mult.inputs[1].default_value = scale
    mult.inputs[2].default_value = translate
    links.new(sep.outputs[channel], mult.inputs[0])
    links.new(mult.outputs[0], combine.inputs["X" if channel == "X" else "Y"])
links.new(combine.outputs[0], projected.inputs["Vector"])
attribute = nodes.new("ShaderNodeAttribute")
attribute.attribute_name = "projection_weight"
mix = nodes.new("ShaderNodeMixRGB")
links.new(attribute.outputs["Fac"], mix.inputs[0])
links.new(original.outputs["Color"], mix.inputs[1])
links.new(projected.outputs["Color"], mix.inputs[2])
links.new(mix.outputs[0], emit.inputs["Color"])
baked = bpy.data.images.new("Astronaut reference atlas", 4096, 4096, alpha=True)
target = nodes.new("ShaderNodeTexImage")
target.image = baked
nodes.active = target
weights = []
bpy.ops.object.select_all(action="DESELECT")
for obj in objects:
    attr = obj.data.attributes.new("projection_weight", "FLOAT", "POINT")
    for v in obj.data.vertices:
        origin = Vector((v.co.x, -2.0, v.co.z))
        hit, normal, index, distance = bvh.ray_cast(origin, Vector((0, 1, 0)))
        visible = hit is not None and abs(hit.y - v.co.y) < 0.004
        facing = max(0.0, min(1.0, (-v.normal.y - 0.15) / 0.5))
        facing = facing * facing * (3 - 2 * facing)
        u = v.co.x * sx + tx
        vv = v.co.z * sz + tz
        px = int(np.clip(u * w, 0, w - 1))
        py = int(np.clip(vv * h, 0, h - 1))
        foreground = pixels[py, px, :3].min() < 0.94
        weight = float(visible and foreground) * facing
        attr.data[v.index].value = weight
        weights.append(weight)
    obj.data.materials.clear()
    obj.data.materials.append(mat)
    obj.select_set(True)
# A joined bake copy writes one atlas; original component objects remain separate.
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 = bpy.context.scene
scene.render.engine = "CYCLES"
scene.cycles.samples = 1
scene.render.bake.margin = 16
scene.render.bake.use_clear = True
bpy.ops.object.bake(type="EMIT")
baked.filepath_raw = str((out / "reference-atlas.png").resolve())
baked.file_format = "PNG"
baked.save()
bpy.data.objects.remove(proxy, do_unlink=True)
final = bpy.data.materials.new("Astronaut painted suit")
final.use_nodes = True
node = final.node_tree.nodes.get("Principled BSDF")
node.inputs["Roughness"].default_value = 0.7
node.inputs["Metallic"].default_value = 0
tex = final.node_tree.nodes.new("ShaderNodeTexImage")
tex.image = baked
final.node_tree.links.new(tex.outputs["Color"], node.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.context.view_layer.objects.active = objects[0]
bpy.ops.export_scene.gltf(
    filepath=str((out / "character-static.glb").resolve()), export_format="GLB", use_selection=True
)
(out / "texture-transfer.json").write_text(
    json.dumps(
        {
            "method": "Visibility- and normal-weighted projection of original reference, baked to existing UVs; original texture retained on unseen surfaces",
            "resolution": 4096,
            "projection": [float(sx), float(sz), float(tx), float(tz)],
            "vertices_with_reference": sum(x > 0.5 for x in weights),
            "vertices": len(weights),
            "roughness": 0.7,
        },
        indent=2,
    )
)