#!/usr/bin/env python3
"""
GLB post-processor for The Settled Reach asset pipeline.

Run via Blender headless:
    flatpak run org.blender.Blender --background --python postprocess_glb.py -- input.glb output.glb [options]

Operations:
    1. Normalize scale to fit a target bounding box (default 1x1x1 world units)
    2. Preserve Trellis texture, bake a recolor mask for the dominant color region
    3. Center the model on the origin, feet on the floor (Y=0)
    4. Export clean .glb + mask PNG sidecar

The mask texture marks the dominant color region (white = replaceable by engine
tint, black = keep original Trellis detail). Godot loads the mask as a second
texture and uses a color-key shader for runtime recoloring.

Options:
    --target-width FLOAT    Target width in world units (default: 1.0)
    --target-height FLOAT   Target height in world units (default: auto from aspect ratio)
    --color-threshold FLOAT Distance threshold for dominant color detection (default: 0.25)
    --material-name NAME    Material slot name (default: mat_primary)
"""

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


def get_script_args():
    """Extract arguments after '--' from Blender's sys.argv."""
    try:
        idx = sys.argv.index("--")
        return sys.argv[idx + 1:]
    except ValueError:
        return []


def parse_args(args):
    """Parse script arguments."""
    if len(args) < 2:
        print("Usage: postprocess_glb.py -- input.glb output.glb [--target-width N] [--color #hex] [--material-name name]")
        sys.exit(1)

    result = {
        "input": args[0],
        "output": args[1],
        "target_width": 1.0,
        "target_height": None,
        "color_threshold": 0.25,
        "material_name": "mat_primary",
    }

    i = 2
    while i < len(args):
        if args[i] == "--target-width" and i + 1 < len(args):
            result["target_width"] = float(args[i + 1])
            i += 2
        elif args[i] == "--target-height" and i + 1 < len(args):
            result["target_height"] = float(args[i + 1])
            i += 2
        elif args[i] == "--color-threshold" and i + 1 < len(args):
            result["color_threshold"] = float(args[i + 1])
            i += 2
        elif args[i] == "--material-name" and i + 1 < len(args):
            result["material_name"] = args[i + 1]
            i += 2
        else:
            print(f"Unknown argument: {args[i]}")
            i += 1

    return result


def clear_scene():
    """Remove all objects from the scene."""
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete()
    # Clear orphan data
    for block in bpy.data.meshes:
        if block.users == 0:
            bpy.data.meshes.remove(block)
    for block in bpy.data.materials:
        if block.users == 0:
            bpy.data.materials.remove(block)
    for block in bpy.data.images:
        if block.users == 0:
            bpy.data.images.remove(block)


def import_glb(filepath):
    """Import a .glb file."""
    print(f"  Importing {filepath}...")
    bpy.ops.import_scene.gltf(filepath=filepath)
    return [obj for obj in bpy.context.scene.objects if obj.type == 'MESH']


def get_combined_bbox(objects):
    """Get the combined bounding box of all mesh objects."""
    min_co = Vector((float('inf'), float('inf'), float('inf')))
    max_co = Vector((float('-inf'), float('-inf'), float('-inf')))

    for obj in objects:
        for corner in obj.bound_box:
            world_co = obj.matrix_world @ Vector(corner)
            min_co.x = min(min_co.x, world_co.x)
            min_co.y = min(min_co.y, world_co.y)
            min_co.z = min(min_co.z, world_co.z)
            max_co.x = max(max_co.x, world_co.x)
            max_co.y = max(max_co.y, world_co.y)
            max_co.z = max(max_co.z, world_co.z)

    return min_co, max_co


def normalize_scale(objects, target_width, target_height=None):
    """Scale all objects so the combined bounding box fits the target size."""
    min_co, max_co = get_combined_bbox(objects)
    size = max_co - min_co

    if size.x == 0 and size.y == 0 and size.z == 0:
        print("  WARNING: Zero-size bounding box, skipping scale normalization")
        return 1.0

    # In Blender: X=right, Y=forward, Z=up
    # Our game: width is max(X, Y), height is Z
    current_width = max(size.x, size.y)
    current_height = size.z

    if current_width == 0:
        current_width = 0.001

    # Scale to target width
    scale_factor = target_width / current_width

    # If target height specified, use the more constraining dimension
    if target_height is not None and current_height > 0:
        height_scale = target_height / current_height
        scale_factor = min(scale_factor, height_scale)

    print(f"  Current size: {size.x:.3f} x {size.y:.3f} x {size.z:.3f}")
    print(f"  Scale factor: {scale_factor:.4f}")
    print(f"  Result size: {size.x * scale_factor:.3f} x {size.y * scale_factor:.3f} x {size.z * scale_factor:.3f}")

    # Apply scale to all objects
    for obj in objects:
        obj.scale *= scale_factor

    # Apply transforms
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.transform_apply(location=False, rotation=False, scale=True)

    return scale_factor


def center_on_origin(objects):
    """Center the combined bounding box on origin, bottom at Z=0."""
    min_co, max_co = get_combined_bbox(objects)
    center = (min_co + max_co) / 2.0

    # Move so center X/Y is at origin, bottom Z is at 0
    offset = Vector((-center.x, -center.y, -min_co.z))

    for obj in objects:
        obj.location += offset

    # Apply location
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.transform_apply(location=True, rotation=False, scale=False)

    print(f"  Centered: offset applied ({offset.x:.3f}, {offset.y:.3f}, {offset.z:.3f})")


def find_texture_image(objects):
    """Find the base color texture image from the imported materials."""
    for obj in objects:
        if obj.type != 'MESH':
            continue
        for slot in obj.material_slots:
            mat = slot.material
            if mat is None or not mat.use_nodes:
                continue
            for node in mat.node_tree.nodes:
                if node.type == 'TEX_IMAGE' and node.image is not None:
                    return node.image
    return None


def analyze_dominant_color(image):
    """Find the dominant color in a Blender image by pixel frequency.

    Quantizes to 8-bit buckets (32 levels per channel) and returns the
    center of the largest bucket as (r, g, b).
    """
    pixels = list(image.pixels)  # flat RGBA
    width, height = image.size
    total = width * height

    # Quantize into buckets (5-bit per channel = 32 levels)
    LEVELS = 32
    buckets = {}
    for i in range(0, len(pixels), 4):
        r, g, b = pixels[i], pixels[i + 1], pixels[i + 2]
        a = pixels[i + 3]
        if a < 0.1:
            continue  # skip transparent
        qr = int(r * (LEVELS - 1))
        qg = int(g * (LEVELS - 1))
        qb = int(b * (LEVELS - 1))
        key = (qr, qg, qb)
        if key in buckets:
            buckets[key][0] += 1
            buckets[key][1] += r
            buckets[key][2] += g
            buckets[key][3] += b
        else:
            buckets[key] = [1, r, g, b]

    if not buckets:
        return (0.8, 0.8, 0.8)

    # Find the largest bucket
    best_key = max(buckets, key=lambda k: buckets[k][0])
    count, sum_r, sum_g, sum_b = buckets[best_key]
    dominant = (sum_r / count, sum_g / count, sum_b / count)

    pct = (count / total) * 100 if total > 0 else 0
    print(f"  Dominant color: ({dominant[0]:.2f}, {dominant[1]:.2f}, {dominant[2]:.2f}) — {pct:.1f}% of pixels")

    return dominant


def generate_mask(image, dominant_color, threshold, output_path):
    """Generate a recolor mask: white where pixels are near the dominant color,
    black elsewhere. Saves as a PNG sidecar next to the GLB."""
    pixels = list(image.pixels)
    width, height = image.size
    dr, dg, db = dominant_color

    mask_pixels = [0.0] * (width * height * 4)
    replaceable_count = 0
    total_count = 0

    for i in range(0, len(pixels), 4):
        px_idx = i // 4
        r, g, b, a = pixels[i], pixels[i + 1], pixels[i + 2], pixels[i + 3]

        if a < 0.1:
            # Transparent — not replaceable
            mask_pixels[i + 3] = 0.0
            continue

        total_count += 1
        # Euclidean distance in RGB space
        dist = ((r - dr) ** 2 + (g - dg) ** 2 + (b - db) ** 2) ** 0.5

        if dist <= threshold:
            # Replaceable region — white
            mask_pixels[i] = 1.0
            mask_pixels[i + 1] = 1.0
            mask_pixels[i + 2] = 1.0
            mask_pixels[i + 3] = 1.0
            replaceable_count += 1
        else:
            # Detail region — black
            mask_pixels[i] = 0.0
            mask_pixels[i + 1] = 0.0
            mask_pixels[i + 2] = 0.0
            mask_pixels[i + 3] = 1.0

    pct = (replaceable_count / total_count * 100) if total_count > 0 else 0
    print(f"  Mask: {replaceable_count}/{total_count} pixels replaceable ({pct:.1f}%)")

    # Create a new Blender image for the mask
    mask_img = bpy.data.images.new("recolor_mask", width, height, alpha=True)
    mask_img.pixels = mask_pixels
    mask_img.filepath_raw = output_path
    mask_img.file_format = 'PNG'
    mask_img.save()
    print(f"  Mask saved: {output_path}")

    return mask_img


def setup_materials(objects, material_name):
    """Keep the original Trellis texture but rename the material slot.
    Sets roughness high and specular low for toon compatibility."""
    for obj in objects:
        if obj.type != 'MESH':
            continue
        for slot in obj.material_slots:
            mat = slot.material
            if mat is None:
                continue
            mat.name = material_name
            if not mat.use_nodes:
                mat.roughness = 1.0
                mat.specular_intensity = 0.0
            else:
                # Find the Principled BSDF and adjust for toon
                for node in mat.node_tree.nodes:
                    if node.type == 'BSDF_PRINCIPLED':
                        node.inputs['Roughness'].default_value = 1.0
                        node.inputs['Specular IOR Level'].default_value = 0.0
    print(f"  Material renamed to '{material_name}', roughness=1.0, specular=0.0")


def export_glb(filepath):
    """Export the scene as .glb, preserving embedded textures."""
    print(f"  Exporting to {filepath}...")
    os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)

    # Ensure all images are packed — Trellis GLBs embed textures, but after
    # Blender import they may become external references that point nowhere.
    packed = 0
    for img in bpy.data.images:
        if img.packed_file is None and img.filepath:
            try:
                img.pack()
                packed += 1
            except Exception as e:
                print(f"  WARNING: Could not pack image '{img.name}': {e}")
    if packed:
        print(f"  Packed {packed} image(s) back into the blend data")

    bpy.ops.export_scene.gltf(
        filepath=filepath,
        export_format='GLB',
        use_selection=False,
        export_apply=True,
        export_materials='EXPORT',
        export_image_format='AUTO',
    )

    size = os.path.getsize(filepath)
    print(f"  Exported: {filepath} ({size} bytes)")


def main():
    args = parse_args(get_script_args())

    print(f"\n=== GLB Post-Processor ===")
    print(f"  Input:  {args['input']}")
    print(f"  Output: {args['output']}")
    print(f"  Target width: {args['target_width']}")
    print(f"  Color threshold: {args['color_threshold']}")
    print(f"  Material: {args['material_name']}")
    print()

    # Clear and import
    clear_scene()
    mesh_objects = import_glb(args['input'])

    if not mesh_objects:
        print("  ERROR: No mesh objects found in .glb")
        sys.exit(1)

    print(f"  Found {len(mesh_objects)} mesh object(s)")

    # Normalize scale
    normalize_scale(mesh_objects, args['target_width'], args['target_height'])

    # Center on origin, feet on floor
    center_on_origin(mesh_objects)

    # Analyze texture and generate recolor mask (if texture exists)
    tex_image = find_texture_image(mesh_objects)
    if tex_image:
        dominant = analyze_dominant_color(tex_image)
        mask_path = os.path.splitext(os.path.abspath(args['output']))[0] + "_mask.png"
        generate_mask(tex_image, dominant, args['color_threshold'], mask_path)
        # Keep original texture, just rename material and adjust PBR
        setup_materials(mesh_objects, args['material_name'])
    else:
        print("  No texture found — Trellis output has no atlas.")
        print("  Keeping existing material as-is (flat color, single surface).")
        setup_materials(mesh_objects, args['material_name'])

    # Export (preserves texture in GLB)
    export_glb(args['output'])

    print("\n=== Done ===\n")


if __name__ == "__main__":
    main()