#!/usr/bin/env python3
"""Bake per-asset ambient occlusion maps for the Blender library.

The GL Compatibility renderer has no SSAO, so the assets carry a baked AO map
instead. This tool opens each generated `.blend` source, bakes AO into a single
image using the asset's existing UVs, and writes it next to the GLB as
`<name>_ao.png`. Godot wires the map into the material's AO channel.

Run inside Blender (flatpak; absolute paths required):

    flatpak run org.blender.Blender --background \
      --python "$PWD/tools/bake_ao.py" -- --only bollard
    flatpak run org.blender.Blender --background \
      --python "$PWD/tools/bake_ao.py"

The tool is a developer step; the committed PNGs mean CI never needs Blender.
"""

import argparse
import json
import sys
from pathlib import Path

try:
    import bpy  # type: ignore
except ImportError:  # pragma: no cover
    bpy = None

ROOT = Path(__file__).resolve().parent.parent
LIBRARY = ROOT / "assets" / "models" / "blender"
SRC = ROOT / "assets" / "models" / "blender_src"
MANIFEST = LIBRARY / "manifest.json"
SIZE = 256
SAMPLES = 16


def bake_asset(name: str, category: str) -> bool:
    source = SRC / f"{name}.blend"
    if not source.is_file():
        print(f"skip {name}: missing {source}")
        return False
    bpy.ops.wm.open_mainfile(filepath=str(source))
    # Some character variants keep tiny rigid parts (shoes/soles) without UVs;
    # bake the UV-mapped body/head/limbs and leave those parts untextured.
    meshes = [o for o in bpy.context.scene.objects
              if o.type == "MESH" and o.data.uv_layers]
    if not meshes:
        print(f"skip {name}: no UV-mapped mesh")
        return False
    image = bpy.data.images.new(f"{name}_ao", width=SIZE, height=SIZE,
                                alpha=False)
    image.generated_color = (1.0, 1.0, 1.0, 1.0)
    for mesh in meshes:
        for material in mesh.data.materials:
            if material is None or not material.use_nodes:
                continue
            nodes = material.node_tree.nodes
            texture = nodes.new("ShaderNodeTexImage")
            texture.image = image
            for node in nodes:
                node.select = False
            texture.select = True
            nodes.active = texture
    for obj in bpy.context.selected_objects:
        obj.select_set(False)
    for mesh in meshes:
        mesh.select_set(True)
    bpy.context.view_layer.objects.active = meshes[0]

    scene = bpy.context.scene
    scene.render.engine = "CYCLES"
    scene.cycles.samples = SAMPLES
    scene.cycles.use_denoising = False
    scene.render.bake.margin = 4
    scene.render.bake.use_clear = True
    bpy.ops.object.bake(type="AO", use_clear=True)

    out = LIBRARY / category / f"{name}_ao.png"
    out.parent.mkdir(parents=True, exist_ok=True)
    image.filepath_raw = str(out)
    image.file_format = "PNG"
    image.save()
    print(f"baked {category}/{name} -> {out.relative_to(ROOT)}")
    return True


def main(argv=None) -> int:
    raw = list(argv if argv is not None else sys.argv[1:])
    if "--" in raw:
        raw = raw[raw.index("--") + 1:]
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--only", default="", help="comma-separated asset names")
    args = parser.parse_args(raw)
    if bpy is None:
        print("bake_ao.py must run inside Blender (flatpak)", file=sys.stderr)
        return 2
    manifest = json.loads(MANIFEST.read_text(encoding="utf-8"))
    only = {name for name in args.only.split(",") if name}
    baked = 0
    for name, entry in sorted(manifest["assets"].items()):
        if only and name not in only:
            continue
        if bake_asset(name, entry["category"]):
            baked += 1
    print(f"baked {baked} AO maps")
    return 0 if baked else 1


if __name__ == "__main__":
    sys.exit(main())