"""
Add missing shape keys from Body_shapes to body part objects.
Preserves existing shape keys.
Usage: blender -b --python add_missing_shape_keys.py -- <blend_file> <output_file>
"""

import bpy
import sys
import os
import mathutils
from mathutils.bvhtree import BVHTree

# Import functions from transfer_shape_keys.py
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from transfer_shape_keys import (
    compute_robust_surface_mapping,
    compute_side_preserving_position,
    smooth_boundary_areas,
    smooth_penetration_areas,
    create_bvh_for_source,
    load_source_data,
    fix_seams_across_objects,
    fix_normals_across_objects
)

def get_source_object_name(target_obj_info):
    if target_obj_info.get('ref_shapes'):
        return target_obj_info['ref_shapes']
    return None

def add_missing_shape_keys(source_data, target_obj):
    """Add only missing shape keys to target object, preserving existing ones."""
    if not target_obj.data.shape_keys:
        target_obj.shape_key_add(name='Basis', from_mix=False)
    
    existing_names = {kb.name for kb in target_obj.data.shape_keys.key_blocks}
    missing_names = [n for n in source_data['names'] if n != 'Basis' and n not in existing_names]
    
    if not missing_names:
        print(f"    {target_obj.name}: all shape keys present, skipping")
        return
    
    print(f"    {target_obj.name}: adding missing keys: {missing_names}")
    
    # Create missing shape keys
    for sk_name in missing_names:
        target_obj.shape_key_add(name=sk_name, from_mix=False)
    
    target_obj.data.shape_keys.use_relative = True
    
    # Compute mapping
    mapping = compute_robust_surface_mapping(target_obj, source_data)
    
    # Set values for missing shape keys
    for sk_name in missing_names:
        print(f"      Setting {sk_name}...")
        sk = target_obj.data.shape_keys.key_blocks[sk_name]
        
        target_obj.data.shape_keys.use_relative = False
        for i, v in enumerate(sk.data):
            if i < len(mapping['target_verts']):
                pos = compute_side_preserving_position(i, sk_name, mapping, source_data, 1.0)
                v.co = pos
        
        sk.data.update()
        target_obj.data.update_tag()
        bpy.context.view_layer.update()
        
        target_obj.data.shape_keys.use_relative = True
        target_obj.data.update_tag()
        bpy.context.view_layer.update()
        
        smooth_boundary_areas(target_obj, sk_name, mapping, source_data)
        smooth_penetration_areas(target_obj, sk_name, mapping, source_data)
        
        print(f"      ✓ Added {sk_name}")

def process_blend(blend_path, output_path):
    print(f"Processing: {blend_path}")
    bpy.ops.wm.open_mainfile(filepath=blend_path)
    
    # Find source object
    source_obj = None
    for name in ['Body_shapes', 'Body']:
        if name in bpy.data.objects and bpy.data.objects[name].type == 'MESH':
            obj = bpy.data.objects[name]
            if obj.data.shape_keys and len(obj.data.shape_keys.key_blocks) > 1:
                source_obj = obj
                print(f"Found source: {name}")
                break
    
    if not source_obj:
        print("ERROR: No source object with shape keys found")
        return False
    
    # Build source data
    source_data = {
        'names': [kb.name for kb in source_obj.data.shape_keys.key_blocks],
        'vertex_positions': {},
        'polygons': [],
        'is_relative': source_obj.data.shape_keys.use_relative,
        'source_obj_name': source_obj.name
    }
    for poly in source_obj.data.polygons:
        source_data['polygons'].append([v for v in poly.vertices])
    for sk in source_obj.data.shape_keys.key_blocks:
        source_data['vertex_positions'][sk.name] = [(v.co.x, v.co.y, v.co.z) for v in sk.data]
    
    print(f"Source shape keys: {source_data['names']}")
    
    # Process body part objects
    required_props = {'age', 'sex', 'slot'}
    modified = False
    for obj in bpy.data.objects:
        if obj.type == 'MESH' and all(p in obj for p in required_props):
            if obj == source_obj:
                continue
            add_missing_shape_keys(source_data, obj)
            modified = True
    
    # Fix seams: synchronize shape key offsets for vertices sharing
    # the same basis position across body parts.
    fix_seams_across_objects()
    fix_normals_across_objects()

    bpy.ops.wm.save_as_mainfile(filepath=output_path)
    print(f"Saved: {output_path}")
    return True

if __name__ == "__main__":
    try:
        args = sys.argv[sys.argv.index("--") + 1:]
        if len(args) >= 2:
            process_blend(args[0], args[1])
        else:
            print("Usage: blender -b -P script.py -- <input.blend> <output.blend>")
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
