import bpy
import bmesh
import sys
import os
import json
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
    
    IMPORTANT: Body file must be loaded FIRST so its objects keep their original names.
    The clothes file may contain reference meshes with the same names as body parts
    (e.g., 'BodyTop' used for weight painting). If clothes are loaded first, the body's
    real objects get renamed to 'BodyTop.001' etc., breaking the ref_part lookup.
    """
    loaded_objects = []
    
    for path in [body_blend_path, clothes_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 mark vertices that are covered.
    Returns a list where 1 = vertex is covered by clothing, 0 = visible.
    Does NOT modify vertex positions -- that used to cause visible steps
    at clothing boundaries because border vertices kept the inward offset
    while exposed neighbours did not."""
    m_body, _, m_body_normal = get_transformation_matrices(target_body)

    num_verts = len(target_body.data.vertices)
    hit_values = [0] * num_verts

    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 (outward) and backward (inward)
        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

    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)")

    # Store custom properties from clothing before they might be lost
    clothing_props = {}
    for prop_name in ["ref_shapes", "ref_ray_length",
                       "has_own_armature", "own_armature_name"]:
        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")

    # Also store metadata tags for export pipeline
    clothing_meta = {}
    for prop_name in ["ref_layer", "ref_part", "garment_id", "tags",
                       "age", "sex", "slot",
                       "has_own_armature", "own_armature_name"]:
        if prop_name in clothing_obj:
            clothing_meta[prop_name] = clothing_obj[prop_name]
            print(f"  Stored meta '{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)
    has_own_arm = clothing_obj.get("has_own_armature", False)
    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
        # Keep hair armatures (objects with own skeleton)
        if not has_own_arm:
            if old_arm not in whitelist:
                bpy.data.objects.remove(old_arm, do_unlink=True)
        else:
            whitelist.add(old_arm)  # preserve hair skeleton
    elif not is_clothing_copy and has_own_arm:
        # Hair has own skeleton but parent is None (or not an armature);
        # find the armature by name and preserve it.
        own_arm_name = clothing_obj.get("own_armature_name", "")
        old_arm = bpy.data.objects.get(own_arm_name)
        if old_arm and old_arm.type == 'ARMATURE' and old_arm not in whitelist:
            whitelist.add(old_arm)
            print(f"  Preserved hair armature '{old_arm.name}' by name for {clothing_name_for_final}")

    # 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 and not clothing_obj.get("has_own_armature", False):
        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

    # Remember clothing's material BEFORE join so we can verify
    # after Blender potentially remaps faces to a stale copy.
    clothing_mat = clothing_obj.active_material
    bpy.ops.object.join()

    # ---- Fix: ensure clothing faces keep the clothing's material ----
    if clothing_mat:
        clothing_slot = None
        stale_slot = None
        for i, slot in enumerate(new_target.material_slots):
            if slot.material == clothing_mat:
                clothing_slot = i
            elif (slot.material and
                  slot.material != clothing_mat and
                  slot.material.name == clothing_mat.name):
                # Same name, different data block -> stale
                stale_slot = i
        if clothing_slot is not None and stale_slot is not None:
            mesh = new_target.data
            fixed = 0
            for poly in mesh.polygons:
                if poly.material_index == stale_slot:
                    poly.material_index = clothing_slot
                    fixed += 1
            if fixed:
                print(f"  Fixed {fixed} faces: stale mat slot "
                      f"{stale_slot} -> clothing slot {clothing_slot}")
    # ----------------------------------------------------------------

    # Copy stored custom properties to combined object
    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")

    # Aggregate clothing metadata onto combined object
    if "_combined_meta" in new_target:
        meta = json.loads(new_target["_combined_meta"])
    else:
        meta = {"garments": [], "tags": [], "layers": []}
    if "garment_id" in clothing_meta:
        meta["garments"].append(clothing_meta["garment_id"])
    else:
        meta["garments"].append(clothing_name_for_final)
    if "tags" in clothing_meta and clothing_meta["tags"]:
        for t in str(clothing_meta["tags"]).split(";"):
            t = t.strip()
            if t and t not in meta["tags"]:
                meta["tags"].append(t)
    if "ref_layer" in clothing_meta:
        meta["layers"].append(int(clothing_meta["ref_layer"]))
    new_target["_combined_meta"] = json.dumps(meta)

    # 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 = []
    skipped_objects = []
    for o in all_objs:
        if o.type != 'MESH' or "ref_layer" in o:
            continue
        if o.name.startswith("cs_") or o.name.startswith("WGT-") or o.name.startswith("MCH-") or o.name.startswith("ORG-"):
            skipped_objects.append(o.name)
            continue
        if len(o.data.vertices) == 0:
            skipped_objects.append(o.name)
            continue
        body_objects.append(o)

    if skipped_objects:
        print(f"Skipped {len(skipped_objects)} non-body mesh objects: {', '.join(skipped_objects[:10])}{'...' if len(skipped_objects) > 10 else ''}")

    required_body_props = {"age", "sex", "slot"}
    valid_body_objects = []
    skipped_body_objects = []
    for body_obj in body_objects:
        missing = [p for p in required_body_props if p not in body_obj]
        if missing:
            print(f"WARNING: Mesh object '{body_obj.name}' is missing required body properties {missing}.")
            print(f"  Treating as helper/reference object (not a body part).")
            print(f"  Available properties: {[k for k in body_obj.keys() if not k.startswith('_')]}")
            skipped_body_objects.append(body_obj.name)
        else:
            valid_body_objects.append(body_obj)

    if skipped_body_objects:
        print(f"Skipped {len(skipped_body_objects)} mesh objects treated as helpers: {', '.join(skipped_body_objects)}")

    body_objects = valid_body_objects

    required_clothing_props = {"ref_layer", "ref_part", "garment_id"}
    clothing_layer1 = []
    clothing_layer2 = []
    skipped_clothing_objects = []
    for o in all_objs:
        if o.type != 'MESH' or "ref_layer" not in o:
            continue
        missing = [p for p in required_clothing_props if p not in o]
        if missing:
            print(f"WARNING: Mesh object '{o.name}' has ref_layer but is missing required clothing properties {missing}.")
            print(f"  Treating as helper/reference object (not clothing).")
            skipped_clothing_objects.append(o.name)
            continue
        if o["ref_layer"] == 1:
            clothing_layer1.append(o)
        elif o["ref_layer"] == 2:
            clothing_layer2.append(o)

    if skipped_clothing_objects:
        print(f"Skipped {len(skipped_clothing_objects)} mesh objects treated as helpers: {', '.join(skipped_clothing_objects)}")

    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")

    body_objects_dict = {obj.name: obj for obj in body_objects}
    whitelist = set()
    combined_objects = []

    # PROCESS LAYER 1
    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

        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}")
        original_clothing_name = clothing_obj.name

        for combined_obj in combined_objects:
            clothing_copy = clothing_obj.copy()
            clothing_copy.data = clothing_obj.data.copy()
            bpy.context.collection.objects.link(clothing_copy)

            for key in clothing_obj.keys():
                clothing_copy[key] = clothing_obj[key]

            clothing_copy["ref_part"] = combined_obj.name

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

            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
    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}")

        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")

    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)}")

    for obj in all_results:
        whitelist.add(obj)
        if "_combined_meta" in obj:
            meta = json.loads(obj["_combined_meta"])
            layers = meta.get("layers", [])
            garments = meta.get("garments", [])
            tags = meta.get("tags", [])

            if layers:
                obj["layer"] = max(layers)
            else:
                obj["layer"] = 0

            if garments:
                obj["garments"] = ";".join(garments)
            else:
                obj["garments"] = ""

            if tags:
                obj["clothing_tags"] = ";".join(sorted(set(tags)))
            else:
                obj["clothing_tags"] = ""

            print(f"Finalized metadata for {obj.name}: layer={obj['layer']}, garments={obj['garments']}, tags={obj['clothing_tags']}")

            del obj["_combined_meta"]

    cleanup_unused_objects(whitelist)

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


if __name__ == "__main__":
    run_batch_combine()
