#!/usr/bin/env python3
# create_glb_colorpie.py — Blender headless STL -> GLB for colorpie
#
# Example:
#   blender-bin-5.0.0 --background --python create_glb_colorpie.py -- \
#     wedge.stl wedge.glb \
#     --hexes FFFFFF,D9D9D9,BFBFBF,808080,606060,404040,202020,000000,FF6600,CC5500
#
# Optional:
#   --roughness 0.6
#   --metallic 0.0
#
# Notes:
#   - Expects one STL containing 10 disconnected wedge solids.
#   - Separates loose parts after import.
#   - Sorts parts by centroid angle around the origin.
#   - Assigns colors 0..9 in that angular order.
#
# 2026-03-06 ChatGPT
# $Header$

import bpy
import bmesh
import sys
import os
import math
from mathutils import Vector


def delete_other_meshes(keep_objs):
    keep_names = {o.name for o in keep_objs}
    victims = [
        o for o in bpy.context.scene.objects
        if o.type == 'MESH' and o.name not in keep_names
    ]

    if victims:
        bpy.ops.object.select_all(action='DESELECT')
        for o in victims:
            o.select_set(True)
        bpy.context.view_layer.objects.active = victims[0]
        bpy.ops.object.delete()

    log(f"Deleted {len(victims)} extra mesh objects; keeping {len(keep_objs)}")

def part_size_score(obj):
    """
    Score a part by bounding-box volume.
    Good enough to distinguish the main wedge from tiny numeral residue.
    """
    d = obj.dimensions
    return d.x * d.y * d.z


def group_parts_by_angle(parts, tolerance_deg=10.0):
    """
    Group parts whose centroid angles are close together.
    For colorpie, each wedge and its tiny residue should land in the same group.
    """
    parts_sorted = sort_parts_by_angle(parts)
    groups = []

    for obj in parts_sorted:
        a = angle_deg_for_obj(obj)
        placed = False

        for g in groups:
            ga = g["angle"]
            diff = abs(a - ga)
            diff = min(diff, 360.0 - diff)
            if diff <= tolerance_deg:
                g["parts"].append(obj)
                # keep a simple running average angle
                g["angle"] = sum(angle_deg_for_obj(p) for p in g["parts"]) / len(g["parts"])
                placed = True
                break

        if not placed:
            groups.append({
                "angle": a,
                "parts": [obj],
            })

    return groups


def pick_slot_pairs(parts, expected_groups=10, tolerance_deg=10.0):
    groups = group_parts_by_angle(parts, tolerance_deg=tolerance_deg)

    log(f"Grouped {len(parts)} raw parts into {len(groups)} angle groups")

    slot_pairs = []

    for gi, g in enumerate(groups):
        g["parts"] = sorted(g["parts"], key=part_size_score, reverse=True)
        log(f"  group {gi:02d}: angle≈{g['angle']:7.3f} count={len(g['parts'])}")
        for pj, obj in enumerate(g["parts"]):
            dims = obj.dimensions
            log(
                f"    part {pj}: name={obj.name!r} "
                f"score={part_size_score(obj):10.3f} "
                f"dims=({dims.x:8.3f},{dims.y:8.3f},{dims.z:8.3f})"
            )

        if len(g["parts"]) < 2:
            die(f"Expected at least 2 parts in group {gi}, found {len(g['parts'])}")

        body = g["parts"][0]
        text = g["parts"][1]

        slot_pairs.append({
            "angle": angle_deg_for_obj(body),
            "body": body,
            "text": text,
        })

    if len(slot_pairs) != expected_groups:
        die(f"Expected {expected_groups} angle groups; found {len(slot_pairs)}")

    slot_pairs = sorted(slot_pairs, key=lambda s: s["angle"])
    return slot_pairs

def pick_main_parts(parts, expected_groups=10, tolerance_deg=10.0):
    groups = group_parts_by_angle(parts, tolerance_deg=tolerance_deg)

    log(f"Grouped {len(parts)} raw parts into {len(groups)} angle groups")

    for gi, g in enumerate(groups):
        g["parts"] = sorted(g["parts"], key=part_size_score, reverse=True)
        log(f"  group {gi:02d}: angle≈{g['angle']:7.3f} count={len(g['parts'])}")
        for pj, obj in enumerate(g["parts"]):
            dims = obj.dimensions
            log(
                f"    part {pj}: name={obj.name!r} "
                f"score={part_size_score(obj):10.3f} "
                f"dims=({dims.x:8.3f},{dims.y:8.3f},{dims.z:8.3f})"
            )

    if len(groups) != expected_groups:
        die(f"Expected {expected_groups} angle groups; found {len(groups)}")

    main_parts = [g["parts"][0] for g in groups]
    return sort_parts_by_angle(main_parts)

def die(msg, rc=2):
    print(f"ERROR: {msg}")
    raise SystemExit(rc)


def log(msg):
    print(msg)


def hex_to_srgb_rgba(hexstr):
    """
    Accepts:
      'CC5500', '#CC5500', 'C50', '#C50', 'CC5500FF'
    Returns:
      (r, g, b, a) in sRGB 0..1 floats
    """
    h = hexstr.strip()
    if h.startswith("#"):
        h = h[1:]
    if len(h) == 3:
        h = "".join([c * 2 for c in h])
    if len(h) == 6:
        h = h + "FF"
    if len(h) != 8:
        die(f"Bad hex color '{hexstr}'. Use RRGGBB or #RRGGBB.")
    r = int(h[0:2], 16) / 255.0
    g = int(h[2:4], 16) / 255.0
    b = int(h[4:6], 16) / 255.0
    a = int(h[6:8], 16) / 255.0
    return (r, g, b, a)


def srgb_chan_to_linear(c):
    # IEC 61966-2-1
    return c / 12.92 if c <= 0.04045 else ((c + 0.055) / 1.055) ** 2.4


def srgb_rgba_to_linear(rgba):
    r, g, b, a = rgba
    return (
        srgb_chan_to_linear(r),
        srgb_chan_to_linear(g),
        srgb_chan_to_linear(b),
        a,
    )


def ensure_material(name, rgba, roughness=0.6, metallic=0.0):
    mat = bpy.data.materials.get(name)
    if not mat:
        mat = bpy.data.materials.new(name=name)
        mat.use_nodes = True

    bsdf = mat.node_tree.nodes.get("Principled BSDF")
    if bsdf:
        bsdf.inputs["Base Color"].default_value = rgba
        bsdf.inputs["Metallic"].default_value = metallic
        bsdf.inputs["Roughness"].default_value = roughness

    return mat


def clear_scene():
    bpy.ops.wm.read_factory_settings(use_empty=True)


def parse_args(argv):
    pos = []
    hexes = None
    roughness = 0.6
    metallic = 0.0

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

        if a == "--hexes":
            if i + 1 >= len(argv):
                die("Missing value after --hexes")
            hexes = [x.strip() for x in argv[i + 1].split(",") if x.strip()]
            i += 2
            continue

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

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

        pos.append(a)
        i += 1

    if len(pos) == 1:
        inp = pos[0]
        base, _ = os.path.splitext(inp)
        outp = base + ".glb"
    elif len(pos) >= 2:
        inp, outp = pos[0], pos[1]
    else:
        die(
            "USAGE: blender --background --python create_glb_colorpie.py -- "
            "input.stl [output.glb] "
            "[--hexes h1,h2,...,h10] [--roughness 0.6] [--metallic 0.0]"
        )

    if not os.path.exists(inp):
        die(f"Input not found: {inp}")

    if hexes is None:
        hexes = [
            "FFFFFF",
            "D9D9D9",
            "BFBFBF",
            "808080",
            "606060",
            "404040",
            "202020",
            "000000",
            "FF6600",
            "CC5500",
        ]

    if len(hexes) != 10:
        die(f"--hexes must contain exactly 10 entries; got {len(hexes)}")

    return inp, outp, hexes, roughness, metallic


def import_stl(inp):
    res = bpy.ops.wm.stl_import(filepath=inp)
    if 'FINISHED' not in res:
        die(f"STL import failed for: {inp}")

    objs = [o for o in bpy.context.scene.objects if o.type == 'MESH']
    if not objs:
        die("No mesh objects after import")
    return objs


def select_only(obj):
    bpy.ops.object.select_all(action='DESELECT')
    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj


def separate_loose_parts_from_single_object(obj):
    """
    Blender STL import often creates one mesh object containing
    many disconnected shells. Separate them into individual objects.
    """
    select_only(obj)

    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='SELECT')
    bpy.ops.mesh.separate(type='LOOSE')
    bpy.ops.object.mode_set(mode='OBJECT')

    parts = [o for o in bpy.context.scene.objects if o.type == 'MESH']
    if not parts:
        die("No mesh objects after Separate by Loose Parts")

    return parts


def world_bbox_center(obj):
    min_v = Vector((1e30, 1e30, 1e30))
    max_v = Vector((-1e30, -1e30, -1e30))

    for corner in obj.bound_box:
        v = obj.matrix_world @ Vector(corner)
        min_v.x = min(min_v.x, v.x)
        min_v.y = min(min_v.y, v.y)
        min_v.z = min(min_v.z, v.z)
        max_v.x = max(max_v.x, v.x)
        max_v.y = max(max_v.y, v.y)
        max_v.z = max(max_v.z, v.z)

    return (min_v + max_v) * 0.5


def angle_deg_for_obj(obj):
    c = world_bbox_center(obj)
    a = math.degrees(math.atan2(c.y, c.x))
    if a < 0:
        a += 360.0
    return a


def sort_parts_by_angle(parts):
    return sorted(parts, key=angle_deg_for_obj)


def assign_material(obj, mat):
    if obj.type != 'MESH':
        return
    if not obj.data.materials:
        obj.data.materials.append(mat)
    else:
        obj.data.materials[0] = mat


def center_all_meshes(objs):
    min_v = Vector((1e30, 1e30, 1e30))
    max_v = Vector((-1e30, -1e30, -1e30))

    for o in objs:
        for corner in o.bound_box:
            v = o.matrix_world @ Vector(corner)
            min_v.x = min(min_v.x, v.x)
            min_v.y = min(min_v.y, v.y)
            min_v.z = min(min_v.z, v.z)
            max_v.x = max(max_v.x, v.x)
            max_v.y = max(max_v.y, v.y)
            max_v.z = max(max_v.z, v.z)

    center = (min_v + max_v) * 0.5

    for o in objs:
        o.location -= center


def export_glb(outp, objs):
    bpy.ops.object.select_all(action='DESELECT')
    for o in objs:
        o.select_set(True)
    bpy.context.view_layer.objects.active = objs[0]

    res = bpy.ops.export_scene.gltf(
        filepath=outp,
        export_format='GLB',
        export_apply=True,
        use_selection=True,
    )
    if 'FINISHED' not in res:
        die(f"GLB export failed: {outp}")

def ensure_text_material(name="PIE_TEXT_BLACK"):
    black = srgb_rgba_to_linear((0.0, 0.0, 0.0, 1.0))
    return ensure_material(name, black, roughness=0.7, metallic=0.0)

def main():
    log("create_glb_colorpie.py starting")
    log(f"Script path: {os.path.abspath(__file__)}")
    log(f"Python executable: {sys.executable}")
    log(f"argv: {sys.argv}")

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

    inp, outp, hexes, roughness, metallic = parse_args(argv)

    clear_scene()

    imported = import_stl(inp)

    if len(imported) == 1:
        parts = separate_loose_parts_from_single_object(imported[0])
    else:
        parts = imported

    slot_pairs = pick_slot_pairs(parts, expected_groups=10, tolerance_deg=10.0)
    text_mat = ensure_text_material()
    export_objs = []

    for i, slot in enumerate(slot_pairs):
        body = slot["body"]
        text = slot["text"]

        hexval = hexes[i]
        srgb = hex_to_srgb_rgba(hexval)
        linear = srgb_rgba_to_linear(srgb)

        body_mat_name = f"PIE_{i}_{hexval.strip().lstrip('#').upper()}"
        body_mat = ensure_material(
            body_mat_name,
            linear,
            roughness=roughness,
            metallic=metallic,
        )

        assign_material(body, body_mat)
        assign_material(text, text_mat)

        body.name = f"wedge_{i}"
        text.name = f"label_{i}"

        export_objs.extend([body, text])

        log(
            f"Assigned wedge_{i}: angle={angle_deg_for_obj(body):7.3f} "
            f"hex=#{hexval.strip().lstrip('#').upper()} "
            f"text={text.name}"
        )

    center_all_meshes(export_objs)
    export_glb(outp, export_objs)
    log(f"Wrote: {outp}")


if __name__ == "__main__":
    main()
