#!/usr/bin/env python3
# 20260308 ChatGPT
# $Header$
#
# Example:
# Note: we need to launch blender in  the same directory as the
# studio_small_03_1k.exr is located, so
#   cd /home/jlpoole/work/colorpie/blender
#   mkdir out
#   blender-bin-5.0.0 -b colorpie.blend \
#     --python ../scripts/render_colorpie_frames.py -- \
#     .. /glb/colorpie_muted_20260308_12303.glb \
#     --outdir out
#
# Optional:
#   --parent ColorpieParent
#   --frame1 1
#   --frame2 144
#   --format PNG
#
# Notes:
#   - This script uses the current .blend as a render template.
#   - It does NOT rebuild lights, world nodes, camera, or compositor.
#   - It deletes current wedge_/label_ descendants under ColorpieParent,
#     imports the replacement GLB, parents imported roots to ColorpieParent,
#     and renders two still frames.

import bpy
import os
import sys
import time


def die(msg, code=1):
    print(f"ERROR: {msg}", file=sys.stderr)
    sys.exit(code)


def log(msg):
    print(msg, flush=True)


def parse_args(argv):
    if "--" in argv:
        argv = argv[argv.index("--") + 1:]
    else:
        argv = []

    glb_path = None
    outdir = "."
    parent_name = "ColorpieParent"
    frame1 = 1
    frame2 = 144
    image_format = "PNG"

    i = 0
    positional = []

    while i < len(argv):
        a = argv[i]

        if a == "--outdir":
            i += 1
            if i >= len(argv):
                die("Missing value after --outdir")
            outdir = argv[i]

        elif a == "--parent":
            i += 1
            if i >= len(argv):
                die("Missing value after --parent")
            parent_name = argv[i]

        elif a == "--frame1":
            i += 1
            if i >= len(argv):
                die("Missing value after --frame1")
            frame1 = int(argv[i])

        elif a == "--frame2":
            i += 1
            if i >= len(argv):
                die("Missing value after --frame2")
            frame2 = int(argv[i])

        elif a == "--format":
            i += 1
            if i >= len(argv):
                die("Missing value after --format")
            image_format = argv[i].upper()

        else:
            positional.append(a)

        i += 1

    if not positional:
        die("Usage: blender -b colorpie.blend --python render_colorpie_frames.py -- input.glb [--outdir DIR]")

    glb_path = positional[0]

    if not os.path.exists(glb_path):
        die(f"GLB not found: {glb_path}")

    os.makedirs(outdir, exist_ok=True)

    return {
        "glb_path": glb_path,
        "outdir": outdir,
        "parent_name": parent_name,
        "frame1": frame1,
        "frame2": frame2,
        "image_format": image_format,
    }


def make_paths_absolute_and_reload_images():
    try:
        bpy.ops.file.make_paths_absolute()
        log("Made external paths absolute")
    except Exception as exc:
        log(f"WARNING: could not make paths absolute: {exc}")

    reloaded = 0
    for img in bpy.data.images:
        try:
            if img.filepath:
                img.reload()
                reloaded += 1
        except Exception as exc:
            log(f"WARNING: could not reload image {img.name}: {exc}")

    log(f"Reloaded {reloaded} image(s)")


def dump_world_info(scene):
    world = scene.world
    log(f"Scene: {scene.name}")
    log(f"Render engine: {scene.render.engine}")

    if world is None:
        log("World: NONE")
        return

    log(f"World: {world.name}")
    log(f"World use_nodes: {world.use_nodes}")

    if world.use_nodes and world.node_tree:
        for node in world.node_tree.nodes:
            if node.type == 'TEX_ENV':
                img = getattr(node, "image", None)
                if img:
                    log(f"Environment texture: {img.name} -> {img.filepath}")
                else:
                    log(f"Environment texture node {node.name} has no image")
    if world.node_tree:
        for node in world.node_tree.nodes:
            log(f"World node: {node.name} type={node.type}")
            if node.type == 'TEX_ENV':
                img = getattr(node, "image", None)
                if img:
                    log(f"Environment texture: {img.name} -> {img.filepath}")
                else:
                    log(f"Environment texture node {node.name} has no image")


def descendant_objects(obj):
    out = []
    stack = list(obj.children)
    while stack:
        cur = stack.pop()
        out.append(cur)
        stack.extend(cur.children)
    return out


def is_colorpie_mesh_name(name):
    return name.startswith("wedge_") or name.startswith("label_")


def delete_old_colorpie_objects(parent):
    victims = []
    for obj in descendant_objects(parent):
        if is_colorpie_mesh_name(obj.name):
            victims.append(obj)

    if not victims:
        log("No existing wedge_/label_ objects found beneath parent")
        return

    # Remove children before parents if any hierarchy exists
    victims = sorted(victims, key=lambda o: len(o.children), reverse=True)

    for obj in victims:
        try:
            bpy.data.objects.remove(obj, do_unlink=True)
        except Exception as exc:
            log(f"WARNING: failed to remove {obj.name}: {exc}")

    log(f"Deleted {len(victims)} old colorpie object(s)")


def import_glb(glb_path):
    before = set(bpy.data.objects.keys())

    res = bpy.ops.import_scene.gltf(filepath=glb_path)
    if "FINISHED" not in res:
        die(f"Import failed for {glb_path}")

    after = set(bpy.data.objects.keys())
    new_names = sorted(after - before)
    imported = [bpy.data.objects[name] for name in new_names]

    if not imported:
        die("Import produced no detectable new objects")

    log(f"Imported {len(imported)} object(s)")
    for obj in imported:
        pname = obj.parent.name if obj.parent else "-"
        log(f"  imported: {obj.name}  type={obj.type}  parent={pname}")

    return imported


def get_import_roots(imported):
    imported_set = set(imported)
    roots = [obj for obj in imported if obj.parent not in imported_set]
    return roots if roots else imported


def parent_imported_to_colorpieparent(imported, parent):
    roots = get_import_roots(imported)

    for obj in roots:
        world_matrix = obj.matrix_world.copy()
        obj.parent = parent
        obj.matrix_parent_inverse = parent.matrix_world.inverted()
        obj.matrix_world = world_matrix

    log(f"Parented {len(roots)} imported root object(s) to {parent.name}")


# def configure_still_output(scene, image_format):
#     fmt = image_format.upper()

#     if fmt == "WEBP":
#         # Blender 5 may expose media_type when scene is in movie mode
#         isettings = scene.render.image_settings

#         if hasattr(isettings, "media_type"):
#             try:
#                 isettings.media_type = 'IMAGE'
#             except Exception as exc:
#                 log(f"WARNING: could not set media_type=IMAGE: {exc}")
#                 log("Falling back to PNG")
#                 fmt = "PNG"

#         if fmt == "WEBP":
#             try:
#                 isettings.file_format = 'WEBP'
#                 if hasattr(isettings, "quality"):
#                     isettings.quality = 95
#                 ext = ".webp"
#                 log("Output format: WEBP")
#                 return ext
#             except Exception as exc:
#                 log(f"WARNING: could not set WEBP output: {exc}")
#                 log("Falling back to PNG")
#                 fmt = "PNG"

#     if fmt == "PNG":
#         scene.render.image_settings.file_format = 'PNG'
#         scene.render.image_settings.color_mode = 'RGBA'
#         ext = ".png"
#         log("Output format: PNG")
#         return ext

#     die(f"Unsupported format requested: {image_format}")

def configure_still_output(scene, image_format):
    # Force the actual render engine you want.
    # Blender 5 uses Eevee Next.
    try:
        scene.render.engine = 'BLENDER_EEVEE_NEXT'
        log("Render engine forced to BLENDER_EEVEE_NEXT")
    except Exception as exc:
        log(f"WARNING: could not set BLENDER_EEVEE_NEXT: {exc}")
        try:
            scene.render.engine = 'BLENDER_EEVEE'
            log("Render engine forced to BLENDER_EEVEE")
        except Exception as exc2:
            die(f"Could not set Eevee render engine: {exc2}")

    isettings = scene.render.image_settings

    # Blender 5.x separates still images from video via media_type.
    if hasattr(isettings, "media_type"):
        try:
            isettings.media_type = 'IMAGE'
            log("media_type set to IMAGE")
        except Exception as exc:
            die(f"Could not set media_type=IMAGE: {exc}")

    fmt = image_format.upper()

    if fmt == "WEBP":
        try:
            isettings.file_format = 'WEBP'
            if hasattr(isettings, "quality"):
                isettings.quality = 95
            ext = ".webp"
            log("Output format: WEBP")
            return ext
        except Exception as exc:
            log(f"WARNING: could not set WEBP output: {exc}")
            log("Falling back to PNG")
            fmt = "PNG"

    if fmt == "PNG":
        try:
            isettings.file_format = 'PNG'
            isettings.color_mode = 'RGBA'
            ext = ".png"
            log("Output format: PNG")
            return ext
        except Exception as exc:
            die(f"Could not set PNG output: {exc}")

    die(f"Unsupported format requested: {image_format}")

def render_frame(scene, frame_no, outfile):
    scene.frame_set(frame_no)
    scene.render.filepath = outfile
    log(f"Rendering frame {frame_no} -> {outfile}")
    bpy.ops.render.render(write_still=True)


def main():
    cfg = parse_args(sys.argv)

    glb_path = cfg["glb_path"]
    outdir = cfg["outdir"]
    parent_name = cfg["parent_name"]
    frame1 = cfg["frame1"]
    frame2 = cfg["frame2"]
    image_format = cfg["image_format"]

    scene = bpy.context.scene
    parent = bpy.data.objects.get(parent_name)
    if parent is None:
        die(f"Parent object not found: {parent_name}")

    make_paths_absolute_and_reload_images()
    dump_world_info(scene)
    log(f"Compositor use_nodes: {scene.use_nodes}")

    delete_old_colorpie_objects(parent)
    imported = import_glb(glb_path)
    parent_imported_to_colorpieparent(imported, parent)

    ext = configure_still_output(scene, image_format)

    stamp = time.strftime("%Y%m%d_%H%M")
    base = os.path.splitext(os.path.basename(glb_path))[0]

    out1 = os.path.join(outdir, f"{base}_frame{frame1:03d}_{stamp}{ext}")
    out2 = os.path.join(outdir, f"{base}_frame{frame2:03d}_{stamp}{ext}")

    render_frame(scene, frame1, out1)
    render_frame(scene, frame2, out2)

    log("Done")
    log(f"Wrote: {out1}")
    log(f"Wrote: {out2}")


if __name__ == "__main__":
    main()