"""Import Poly Haven island_tree_01 and export decimated GLBs.

Each output GLB contains exactly one tree variant (1 branch mesh + 1 leaf
mesh) with shared material references so the main generator can instance them
the same way it instances grass tufts.
"""

import math
import os
import sys

import bpy


TREE_BLEND = os.path.join(
    os.path.dirname(__file__), "..", "..",
    "assets", "models", "polyhaven", "island_tree_01",
    "island_tree_01_1k.blend",
)

TARGET_BRANCH_TRIS = 180
TARGET_LEAF_TRIS = 220

PAIRS = [("a", "b"), ("b", "a"), ("c", "a"), ("d", "b")]


def cli_args():
    values = {"out": None}
    argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    i = 0
    while i < len(argv):
        if argv[i].startswith("--") and i + 1 < len(argv):
            values[argv[i][2:]] = argv[i + 1]
            i += 2
        else:
            i += 1
    if not values.get("out"):
        raise RuntimeError("--out is required")
    return values


def decimate_to_target(obj, target_tris):
    source = obj.data
    source.calc_loop_triangles()
    source_tris = len(source.loop_triangles)
    if source_tris <= target_tris:
        return source.copy()
    modifier = obj.modifiers.new("TreeDecimate", "DECIMATE")
    modifier.ratio = max(0.05, target_tris / source_tris)
    bpy.context.view_layer.update()
    depsgraph = bpy.context.evaluated_depsgraph_get()
    evaluated = obj.evaluated_get(depsgraph)
    result = bpy.data.meshes.new_from_object(evaluated)
    result.calc_loop_triangles()
    return result


def export_one(out_path, b_mesh, l_mesh, materials):
    """Export a single tree variant to a clean GLB.

    Works in a fresh, empty scene so nothing from the source blend leaks in.
    """
    scene = bpy.data.scenes.new("_TreeIngest")
    bpy.context.window.scene = scene
    col = bpy.data.collections.new("_Tree")
    scene.collection.children.link(col)

    b_obj = bpy.data.objects.new("Branch", b_mesh)
    l_obj = bpy.data.objects.new("Leaf", l_mesh)
    col.objects.link(b_obj)
    col.objects.link(l_obj)

    b_obj.data.materials.append(materials["branch"])
    l_obj.data.materials.append(materials["leaf"])
    l_obj.data.materials.append(materials["branch"])

    scene.view_layers[0].update()

    bpy.ops.export_scene.gltf(
        filepath=out_path,
        export_format="GLB",
        use_selection=False,
        export_apply=False,
        export_texcoords=True,
        export_normals=True,
        export_materials="EXPORT",
        export_image_format="JPEG",
    )

    bpy.data.scenes.remove(scene, do_unlink=True)


def main():
    args = cli_args()
    out_dir = args["out"]
    os.makedirs(out_dir, exist_ok=True)

    if not os.path.exists(TREE_BLEND):
        print(f"Missing: {TREE_BLEND}")
        return 1

    bpy.ops.wm.open_mainfile(filepath=TREE_BLEND)

    branch_mat = bpy.data.materials.get("island_tree_01_branches")
    leaf_mat = bpy.data.materials.get("island_tree_01_leaves")
    if not branch_mat or not leaf_mat:
        return 1

    decimated = []
    for vi, (branch_id, leaf_id) in enumerate(PAIRS):
        branch_obj = bpy.data.objects.get(
            f"island_tree_01_branches_{branch_id}_LOD1")
        leaf_obj = bpy.data.objects.get(
            f"island_tree_01_leaves_{leaf_id}_LOD1")
        if not branch_obj or not leaf_obj:
            continue

        b_mesh = decimate_to_target(branch_obj, TARGET_BRANCH_TRIS)
        l_mesh = decimate_to_target(leaf_obj, TARGET_LEAF_TRIS)
        decimated.extend([b_mesh, l_mesh])

        out_path = os.path.join(out_dir, f"tree_variant_{vi:02d}.glb")
        export_one(out_path, b_mesh, l_mesh,
                   {"branch": branch_mat, "leaf": leaf_mat})
        kb = os.path.getsize(out_path) // 1024
        print(f"  variant_{vi:02d}: branch={len(b_mesh.loop_triangles)}tris "
              f"leaf={len(l_mesh.loop_triangles)}tris ({kb}KB)")

    for mesh in decimated:
        if mesh and mesh.name in bpy.data.meshes:
            bpy.data.meshes.remove(mesh, do_unlink=True)

    print(f"\nDone — {len(PAIRS)} variants → {out_dir}")
    return 0


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