#!/usr/bin/env python3
# 20260228 ChatGPT
# $Header$
# $Id$
#
# SCRIPT_ID: render_glb_variants.py
# SCRIPT_VERSION: 6
#
# Goal:
#   Headless-safe PNG rendering for GLB models in environments where EGL/OpenGL may crash.
#   Uses Cycles CPU (no Eevee/Workbench).
#
# Usage:
#   blender -b -P blender/render_glb_variants.py -- /path/to/part.glb /path/to/out_prefix

import bpy
import sys
import math
from mathutils import Vector

SCRIPT_ID = "render_glb_variants.py"
SCRIPT_VERSION = 6

def log(msg):
    print(f"{SCRIPT_ID} v{SCRIPT_VERSION}: {msg}")

def die(msg):
    print(f"{SCRIPT_ID} v{SCRIPT_VERSION} ERROR: {msg}", file=sys.stderr)
    raise SystemExit(1)

def argv_after_double_dash():
    if "--" not in sys.argv:
        return []
    return sys.argv[sys.argv.index("--") + 1:]

def reset_scene():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    scene = bpy.context.scene

    scene.render.engine = 'CYCLES'
    # Force CPU
    scene.cycles.device = 'CPU'
    # Fast-ish defaults for documentation renders
    scene.cycles.samples = 32
    scene.cycles.preview_samples = 16
    scene.cycles.use_denoising = True

    scene.render.resolution_x = 1200
    scene.render.resolution_y = 1200
    scene.render.image_settings.file_format = 'PNG'
    scene.render.film_transparent = False
    scene.render.use_freestyle = False

    return scene

def set_world_background_white(scene, strength=1.0):
    if scene.world is None:
        scene.world = bpy.data.worlds.new("World")
    world = scene.world
    world.use_nodes = True
    nt = world.node_tree
    nt.nodes.clear()
    out = nt.nodes.new("ShaderNodeOutputWorld")
    bg  = nt.nodes.new("ShaderNodeBackground")
    bg.inputs["Color"].default_value = (1.0, 1.0, 1.0, 1.0)
    bg.inputs["Strength"].default_value = strength
    nt.links.new(bg.outputs["Background"], out.inputs["Surface"])

def import_glb(path):
    before = set(bpy.data.objects)
    bpy.ops.import_scene.gltf(filepath=path)
    after = set(bpy.data.objects)
    new_objs = [o for o in (after - before)]
    meshes = [o for o in new_objs if o.type == 'MESH']
    if not meshes:
        meshes = [o for o in bpy.context.scene.objects if o.type == 'MESH']
    if not meshes:
        die(f"Failed to import GLB (no mesh objects found): {path}")

    if len(meshes) > 1:
        bpy.ops.object.select_all(action='DESELECT')
        for o in meshes:
            o.select_set(True)
        bpy.context.view_layer.objects.active = meshes[0]
        bpy.ops.object.join()
        obj = bpy.context.view_layer.objects.active
    else:
        obj = meshes[0]

    obj.name = "PART"
    return obj

def center_origin_and_location(obj):
    bpy.ops.object.select_all(action='DESELECT')
    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj
    bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY', center='BOUNDS')
    obj.location = (0, 0, 0)

def bounds_world(obj):
    mw = obj.matrix_world
    corners = [mw @ Vector(c) for c in obj.bound_box]
    xs = [c.x for c in corners]; ys = [c.y for c in corners]; zs = [c.z for c in corners]
    min_v = Vector((min(xs), min(ys), min(zs)))
    max_v = Vector((max(xs), max(ys), max(zs)))
    size = max_v - min_v
    center = (min_v + max_v) / 2.0
    return center, size

def ensure_camera(scene):
    cam_data = bpy.data.cameras.new("OrthoCam")
    cam_data.type = 'ORTHO'
    cam = bpy.data.objects.new("OrthoCam", cam_data)
    scene.collection.objects.link(cam)
    scene.camera = cam
    return cam

def ensure_light_rig(scene):
    # Cycles likes real lights
    def add_area(name, loc, power):
        ld = bpy.data.lights.new(name, type='AREA')
        ld.energy = power
        lo = bpy.data.objects.new(name, ld)
        lo.location = loc
        scene.collection.objects.link(lo)
    add_area("Key", (3, -3, 3), 800)
    add_area("Fill", (-3, 3, 2), 400)
    add_area("Top", (0, 0, 4), 300)

def set_isometric_camera(cam):
    cam.location = (3.0, -3.0, 3.0)
    cam.rotation_euler = (math.radians(55), 0, math.radians(45))

def fit_camera_to_object(cam, obj, margin=1.20):
    _, size = bounds_world(obj)
    max_dim = max(size.x, size.y, size.z)
    cam.data.ortho_scale = max_dim * margin

def set_input_if_exists(node, socket_name, value):
    if socket_name in node.inputs:
        node.inputs[socket_name].default_value = value
        return True
    return False

def make_material_clay(name="Clay", color=(0.83, 0.83, 0.83, 1.0)):
    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    nt = mat.node_tree
    nt.nodes.clear()

    out = nt.nodes.new("ShaderNodeOutputMaterial")
    bsdf = nt.nodes.new("ShaderNodeBsdfPrincipled")

    set_input_if_exists(bsdf, "Base Color", color)
    set_input_if_exists(bsdf, "Roughness", 1.0)
    if not set_input_if_exists(bsdf, "Specular", 0.0):
        set_input_if_exists(bsdf, "Specular IOR Level", 0.0)

    nt.links.new(bsdf.outputs["BSDF"], out.inputs["Surface"])
    return mat

def make_material_black_emission(name="WireBlack"):
    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    nt = mat.node_tree
    nt.nodes.clear()

    out = nt.nodes.new("ShaderNodeOutputMaterial")
    em  = nt.nodes.new("ShaderNodeEmission")
    em.inputs["Color"].default_value = (0, 0, 0, 1)
    em.inputs["Strength"].default_value = 1.0
    nt.links.new(em.outputs["Emission"], out.inputs["Surface"])
    return mat

def make_material_feature_ao(name="FeatureAO"):
    """
    Cavity-like look using Ambient Occlusion node when available.
    Falls back to normal-vector shading if AO node isn't available.
    """
    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    nt = mat.node_tree
    nt.nodes.clear()

    out = nt.nodes.new("ShaderNodeOutputMaterial")
    bsdf = nt.nodes.new("ShaderNodeBsdfPrincipled")
    set_input_if_exists(bsdf, "Roughness", 1.0)
    if not set_input_if_exists(bsdf, "Specular", 0.0):
        set_input_if_exists(bsdf, "Specular IOR Level", 0.0)

    # Try AO node
    try:
        ao = nt.nodes.new("ShaderNodeAmbientOcclusion")
        ao.inputs["Distance"].default_value = 0.2
        ramp = nt.nodes.new("ShaderNodeValToRGB")
        # Slight contrast boost
        ramp.color_ramp.elements[0].position = 0.35
        ramp.color_ramp.elements[1].position = 0.85
        nt.links.new(ao.outputs["AO"], ramp.inputs["Fac"])
        nt.links.new(ramp.outputs["Color"], bsdf.inputs["Base Color"])
    except Exception:
        # Fallback: normals shading
        geo = nt.nodes.new("ShaderNodeNewGeometry")
        add = nt.nodes.new("ShaderNodeVectorMath"); add.operation = 'ADD'
        add.inputs[1].default_value = (1.0, 1.0, 1.0)
        mul = nt.nodes.new("ShaderNodeVectorMath"); mul.operation = 'MULTIPLY'
        mul.inputs[1].default_value = (0.5, 0.5, 0.5)
        nt.links.new(geo.outputs["Normal"], add.inputs[0])
        nt.links.new(add.outputs["Vector"], mul.inputs[0])
        nt.links.new(mul.outputs["Vector"], bsdf.inputs["Base Color"])

    nt.links.new(bsdf.outputs["BSDF"], out.inputs["Surface"])
    return mat

def assign_single_material(obj, mat):
    obj.data.materials.clear()
    obj.data.materials.append(mat)

def duplicate_object(obj, name):
    dup = obj.copy()
    dup.data = obj.data.copy()
    dup.name = name
    bpy.context.scene.collection.objects.link(dup)
    return dup

def delete_object(obj):
    bpy.ops.object.select_all(action='DESELECT')
    obj.select_set(True)
    bpy.ops.object.delete()

def enable_freestyle(scene, thickness=1.0):
    scene.render.use_freestyle = True
    vl = scene.view_layers[0]
    fs = vl.freestyle_settings
    if not fs.linesets:
        fs.linesets.new("LineSet")
    ls = fs.linesets[0]
    try:
        ls.linestyle.thickness = thickness
    except Exception:
        pass

def render(scene, outpath):
    scene.render.filepath = outpath
    bpy.ops.render.render(write_still=True)

def main():
    log("starting")
    args = argv_after_double_dash()
    if len(args) != 2:
        die("Expected: <input.glb> <out_prefix>")

    in_glb, out_prefix = args
    scene = reset_scene()
    set_world_background_white(scene)

    obj = import_glb(in_glb)
    center_origin_and_location(obj)

    ensure_light_rig(scene)
    cam = ensure_camera(scene)
    set_isometric_camera(cam)
    fit_camera_to_object(cam, obj, margin=1.20)

    clay = make_material_clay()
    wire_black = make_material_black_emission()
    feature = make_material_feature_ao()

    # 1) SOLID
    scene.render.use_freestyle = False
    assign_single_material(obj, clay)
    render(scene, f"{out_prefix}_solid.png")

    # 2) SOLID + EDGES
    scene.render.use_freestyle = False
    assign_single_material(obj, clay)
    enable_freestyle(scene, thickness=1.0)
    render(scene, f"{out_prefix}_solid_edges.png")

    # 3) PURE WIRE
    scene.render.use_freestyle = False
    wire_obj = duplicate_object(obj, "PART_WIRE")
    obj.hide_render = True
    wire_obj.hide_render = False
    assign_single_material(wire_obj, wire_black)
    mod = wire_obj.modifiers.new("Wireframe", type='WIREFRAME')
    mod.thickness = 0.0015
    mod.use_replace = True
    mod.use_boundary = True
    render(scene, f"{out_prefix}_wire.png")
    obj.hide_render = False
    delete_object(wire_obj)

    # 4) FEATURE POP (AO/cavity-ish)
    scene.render.use_freestyle = False
    assign_single_material(obj, feature)
    render(scene, f"{out_prefix}_feature.png")

    # 5) FREESTYLE ONLY
    scene.render.use_freestyle = False
    white = make_material_clay(name="WhiteClay", color=(1.0, 1.0, 1.0, 1.0))
    assign_single_material(obj, white)
    enable_freestyle(scene, thickness=1.2)
    render(scene, f"{out_prefix}_freestyle.png")

    log("done")

if __name__ == "__main__":
    main()