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

def load_blend_files(clothes_blend_path, body_blend_path):
    """Load objects from blend files and return all loaded objects"""
    loaded_objects = []
    
    for path in [clothes_blend_path, body_blend_path]:
        with bpy.data.libraries.load(path) as (data_from, data_to):
            data_to.objects = data_from.objects
        for obj in data_to.objects:
            if obj:
                bpy.context.collection.objects.link(obj)
                loaded_objects.append(obj)
    
    return loaded_objects

def setup_bvh_and_matrices(obj):
    """Setup BVH tree and transformation matrices for an object"""
    depsgraph = bpy.context.evaluated_depsgraph_get()
    obj_eval = obj.evaluated_get(depsgraph)
    bvh = BVHTree.FromObject(obj_eval, depsgraph)
    
    return bvh

def get_transformation_matrices(obj):
    """Get transformation matrices for an object"""
    m = obj.matrix_world
    m_inv = m.inverted()
    m_normal = m.to_3x3().inverted().transposed()
    
    return m, m_inv, m_normal

def raycast_and_adjust_vertices(target_body, bvh_cloth, out_ray_length=0.15):
    """Raycast from body to cloth and adjust vertices that intersect"""
    m_body, m_body_inv, m_body_normal = get_transformation_matrices(target_body)
    
    num_verts = len(target_body.data.vertices)
    hit_values = [0] * num_verts
    has_shape_keys = target_body.data.shape_keys is not None
    
    # Forward raycast (into cloth)
    for i, v in enumerate(target_body.data.vertices):
        v_world = m_body @ v.co
        n_world = (m_body_normal @ v.normal).normalized()
        
        # Raycast forward (into cloth) and backward (from inside cloth)
        hit_f, _, _, _ = bvh_cloth.ray_cast(v_world, n_world, out_ray_length)
        hit_b, _, _, _ = bvh_cloth.ray_cast(v_world, -n_world, 0.005)
        
        if hit_f or hit_b:
            hit_values[i] = 1
            # Adjust vertex position to be slightly outside cloth
            offset = -n_world * (0.005 if hit_f else 0.01)
            new_co = m_body_inv @ (v_world + offset)
            v.co = new_co
            
            # Update shape keys if they exist
            if has_shape_keys:
                for kb in target_body.data.shape_keys.key_blocks:
                    kb.data[i].co = new_co
    
    return hit_values

def protect_and_remove_hidden_geometry(target_body, hit_values, threshold=4.0):
    """Protect visible vertices and remove hidden geometry"""
    bm = bmesh.new()
    bm.from_mesh(target_body.data)
    bm.verts.ensure_lookup_table()
    
    # Phase 1: Identify "Layer 1 Border" (Immediate neighbors of visible verts)
    border_l1 = set()
    for v in bm.verts:
        if hit_values[v.index] == 0:  # Visible vertex
            for edge in v.link_edges:
                neighbor = edge.other_vert(v)
                if hit_values[neighbor.index] == 1:
                    border_l1.add(neighbor.index)
    
    # Phase 2: Identify "Layer 2 Buffer" (Neighbors of Layer 1)
    border_l2 = set()
    for idx in border_l1:
        v = bm.verts[idx]
        for edge in v.link_edges:
            neighbor = edge.other_vert(v)
            if hit_values[neighbor.index] == 1:
                border_l2.add(neighbor.index)
    
    # Merge all protected vertices
    protected_indices = set(border_l1) | set(border_l2)
    for i, val in enumerate(hit_values):
        if val == 0:  # Visible vertices
            protected_indices.add(i)
    
    # Deletion logic
    to_delete = []
    for v in bm.verts:
        if v.index in protected_indices:
            continue
        
        # Sum hits of neighbors
        neighbor_hit_sum = hit_values[v.index]
        for edge in v.link_edges:
            neighbor = edge.other_vert(v)
            neighbor_hit_sum += hit_values[neighbor.index]
        
        if neighbor_hit_sum >= threshold:
            to_delete.append(v)
        elif len(v.link_edges) == 1:  # Loose vertices
            to_delete.append(v)
    
    # Perform deletion
    bmesh.ops.delete(bm, geom=to_delete, context='VERTS')
    bm.to_mesh(target_body.data)
    bm.free()
    target_body.data.update()

def process_clothing_pair(clothing_obj, target_obj, whitelist, is_clothing_copy=False, original_clothing_name=None):
    """Process a clothing-body pair

    Args:
        clothing_obj: The clothing object to process
        target_obj: The target object to combine with (body or combined object)
        whitelist: Set of objects to keep
        is_clothing_copy: Whether clothing_obj is a copy (for layer 2 processing)
        original_clothing_name: Original name of clothing if it's a copy (for layer 2 naming)
    """
    # Create a copy of the target object
    new_target = target_obj.copy()
    new_target.data = target_obj.data.copy()
    bpy.context.collection.objects.link(new_target)

    # Copy custom properties
    for key in target_obj.keys():
        new_target[key] = target_obj[key]

    # Ensure the copy has the same transformations
    new_target.matrix_world = target_obj.matrix_world.copy()

    target_name = target_obj.name

    # Determine the name to use for the clothing in the final combined object
    if is_clothing_copy and original_clothing_name:
        clothing_name_for_final = original_clothing_name
    else:
        clothing_name_for_final = clothing_obj.name

    print(f"Processing: {clothing_name_for_final} -> {target_name} (using copy)")

    # NEW CODE: Store custom properties from clothing before they might be lost
    clothing_props = {}
    for prop_name in ["ref_shapes", "ref_ray_length"]:
        if prop_name in clothing_obj:
            clothing_props[prop_name] = clothing_obj[prop_name]
            print(f"  Stored property '{prop_name}' = {clothing_obj[prop_name]} from clothing")

    # Step A: Raycast & adjust vertices
    out_ray_length = 0.015
    if "ref_ray_length" in clothing_obj:
        out_ray_length = clothing_obj["ref_ray_length"]
    bvh_cloth = setup_bvh_and_matrices(clothing_obj)
    hit_values = raycast_and_adjust_vertices(new_target, bvh_cloth, out_ray_length)

    # Step B: Remove hidden geometry
    protect_and_remove_hidden_geometry(new_target, hit_values)

    # Step C: Handle armature and join
    master_arm = new_target.parent if (new_target.parent and new_target.parent.type == 'ARMATURE') else None
    if master_arm:
        whitelist.add(master_arm)

    # Handle clothing armature (if it's not a copy for layer 2)
    if not is_clothing_copy and clothing_obj.parent and clothing_obj.parent.type == 'ARMATURE':
        old_arm = clothing_obj.parent
        clothing_obj.matrix_world = clothing_obj.matrix_world.copy()
        clothing_obj.parent = None
        if old_arm not in whitelist:
            bpy.data.objects.remove(old_arm, do_unlink=True)

    # For layer 2 clothing copies, we don't need to handle armature separately
    # as they'll inherit from the target

    # Reparent to master armature if exists
    if master_arm:
        clothing_obj.parent = master_arm
        for mod in clothing_obj.modifiers:
            if mod.type == 'ARMATURE':
                mod.object = master_arm

    # Join clothing with target copy
    bpy.ops.object.select_all(action='DESELECT')
    clothing_obj.select_set(True)
    new_target.select_set(True)
    bpy.context.view_layer.objects.active = new_target
    bpy.ops.object.join()

    # NEW CODE: Copy stored custom properties to combined object
    # After joining, new_target is the active object and contains the combined mesh
    for prop_name, prop_value in clothing_props.items():
        new_target[prop_name] = prop_value
        print(f"  Copied property '{prop_name}' = {prop_value} to combined object")

    # Rename the combined object using the appropriate clothing name
    new_target.name = f"{target_name}_{clothing_name_for_final}"
    whitelist.add(new_target)

    return new_target

def cleanup_unused_objects(whitelist):
    """Remove all objects not in whitelist"""
    for obj in bpy.data.objects[:]:
        if obj.type in {'MESH', 'ARMATURE'} and obj not in whitelist:
            bpy.data.objects.remove(obj, do_unlink=True)
    
    bpy.ops.outliner.orphans_purge(do_local_ids=True, do_linked_ids=True, do_recursive=True)

def run_batch_combine():
    """Main function to run the batch combine process"""
    try:
        args = sys.argv[sys.argv.index("--") + 1:]
        body_blend_path, clothes_blend_path, output_path = args[0], args[1], args[2]
    except (ValueError, IndexError):
        print("Usage: blender -b -P script.py -- <body.blend> <clothes.blend> <output.blend>")
        return

    # Start fresh
    bpy.ops.wm.read_homefile(use_empty=True)

    # Load objects from blend files
    loaded_objects = load_blend_files(clothes_blend_path, body_blend_path)

    # Categorize loaded objects
    all_objs = bpy.data.objects

    # Separate body objects (no ref_layer property)
    body_objects = [o for o in all_objs if o.type == 'MESH' and "ref_layer" not in o]

    # Separate clothing by layer
    clothing_layer1 = [o for o in all_objs if o.type == 'MESH' and "ref_layer" in o and o["ref_layer"] == 1]
    clothing_layer2 = [o for o in all_objs if o.type == 'MESH' and "ref_layer" in o and o["ref_layer"] == 2]

    print(f"Found {len(body_objects)} body objects")
    print(f"Found {len(clothing_layer1)} layer 1 clothing objects")
    print(f"Found {len(clothing_layer2)} layer 2 clothing objects")

    # Create dictionary of body objects for quick lookup
    body_objects_dict = {obj.name: obj for obj in body_objects}

    # Track objects to keep
    whitelist = set()

    # List to store combined objects from layer 1
    combined_objects = []

    # PROCESS LAYER 1: Combine with original body parts
    print("\n=== PROCESSING LAYER 1 CLOTHING ===")
    for clothing_obj in clothing_layer1:
        if "ref_part" not in clothing_obj:
            print(f"Warning: Layer 1 clothing '{clothing_obj.name}' missing ref_part property, skipping")
            continue

        target_name = clothing_obj["ref_part"]
        original_body = body_objects_dict.get(target_name)

        if not original_body or original_body.type != 'MESH':
            print(f"Warning: Target body '{target_name}' not found for clothing '{clothing_obj.name}', skipping")
            continue

        # Process the pair
        result = process_clothing_pair(clothing_obj, original_body, whitelist)
        if result:
            combined_objects.append(result)
            print(f"Added '{result.name}' to combined objects list")

    print(f"Layer 1 complete. Created {len(combined_objects)} combined objects")

    # PROCESS LAYER 2 (First Pass): Combine with layer 1 results
    print("\n=== PROCESSING LAYER 2 CLOTHING (First Pass - Over Layer 1) ===")
    layer2_results_over_l1 = []

    for clothing_obj in clothing_layer2:
        print(f"\nProcessing layer 2 clothing: {clothing_obj.name}")

        # Store the original clothing name for later use
        original_clothing_name = clothing_obj.name

        # For each combined object from layer 1
        for combined_obj in combined_objects:
            # Create a copy of the layer 2 clothing
            clothing_copy = clothing_obj.copy()
            clothing_copy.data = clothing_obj.data.copy()
            bpy.context.collection.objects.link(clothing_copy)

            # Copy custom properties
            for key in clothing_obj.keys():
                clothing_copy[key] = clothing_obj[key]

            # Set the ref_part to point to the combined object
            clothing_copy["ref_part"] = combined_obj.name

            print(f"  Combining with layer 1 result: {combined_obj.name}")

            # Process the pair
            result = process_clothing_pair(clothing_copy, combined_obj, whitelist,
                                          is_clothing_copy=True,
                                          original_clothing_name=original_clothing_name)
            if result:
                layer2_results_over_l1.append(result)
                print(f"  Created: {result.name}")

    print(f"Layer 2 first pass complete. Created {len(layer2_results_over_l1)} combined objects")

    # PROCESS LAYER 2 (Second Pass): Combine directly with body parts (like layer 1)
    print("\n=== PROCESSING LAYER 2 CLOTHING (Second Pass - Direct to Body) ===")
    layer2_results_direct = []

    for clothing_obj in clothing_layer2:
        if "ref_part" not in clothing_obj:
            print(f"Warning: Layer 2 clothing '{clothing_obj.name}' missing ref_part property, skipping")
            continue

        target_name = clothing_obj["ref_part"]
        original_body = body_objects_dict.get(target_name)

        if not original_body or original_body.type != 'MESH':
            print(f"Warning: Target body '{target_name}' not found for clothing '{clothing_obj.name}', skipping")
            continue

        print(f"\nProcessing layer 2 clothing directly with body: {clothing_obj.name} -> {target_name}")

        # Process directly with body part (no copy needed for the clothing itself)
        result = process_clothing_pair(clothing_obj, original_body, whitelist, is_clothing_copy=False)
        if result:
            layer2_results_direct.append(result)
            print(f"  Created: {result.name}")

    print(f"Layer 2 second pass complete. Created {len(layer2_results_direct)} combined objects")

    # Add all results to combined objects list
    all_results = combined_objects + layer2_results_over_l1 + layer2_results_direct

    print(f"\n=== SUMMARY ===")
    print(f"Layer 1 results: {len(combined_objects)}")
    print(f"Layer 2 over layer 1 results: {len(layer2_results_over_l1)}")
    print(f"Layer 2 direct to body results: {len(layer2_results_direct)}")
    print(f"Total combined objects: {len(all_results)}")

    # Final cleanup - keep all combined objects and their armatures
    for obj in all_results:
        whitelist.add(obj)

    cleanup_unused_objects(whitelist)

    # Save the result
    bpy.ops.wm.save_as_mainfile(filepath=output_path)
    print(f"\nSaved to: {output_path}")


if __name__ == "__main__":
    run_batch_combine()

