import bpy
import re
import os

def link_bsdf_materials():
    """Link all materials from the BSDF library file"""
    
    library_path = r"A:\1 Amazon_Active_Projects\1 BlenderAssets\Amazon\MATERIALS_BSDF_pallette_v1.0.blend"
    
    if not os.path.exists(library_path):
        print(f"Warning: Library file not found at {library_path}")
        return []
    
    print(f"Linking materials from: {library_path}")
    
    # Get list of materials before linking
    materials_before = set(bpy.data.materials.keys())
    
    # Link all materials from the library file
    with bpy.data.libraries.load(library_path, link=True) as (data_from, data_to):
        # Link all materials
        data_to.materials = data_from.materials
    
    # Get list of newly linked materials
    materials_after = set(bpy.data.materials.keys())
    newly_linked = materials_after - materials_before
    
    print(f"Linked {len(newly_linked)} materials from library")
    for mat_name in sorted(newly_linked):
        print(f"  - {mat_name}")
    
    return list(newly_linked)

def remap_appended_to_linked():
    """Remap any appended BSDF materials to their linked counterparts"""
    
    print("\nChecking for appended BSDF materials to remap to linked versions...")
    
    materials = bpy.data.materials
    remapping_count = 0
    
    # Group materials by base name (without library suffix)
    material_groups = {}
    
    for mat in materials:
        # Check if it's a BSDF material (from any source)
        if mat.name.startswith("BSDF_") or "BSDF_" in mat.name:
            # Extract base name (remove library reference if present)
            base_name = mat.name.split(".blend")[0] if ".blend" in mat.name else mat.name
            base_name = base_name.split(".")[0] if "." in base_name else base_name
            
            if base_name not in material_groups:
                material_groups[base_name] = []
            material_groups[base_name].append(mat)
    
    # For each group, prefer linked materials over appended ones
    for base_name, mats in material_groups.items():
        if len(mats) > 1:
            # Sort to prefer linked materials (they have library references)
            linked_mats = [m for m in mats if m.library is not None]
            appended_mats = [m for m in mats if m.library is None]
            
            if linked_mats and appended_mats:
                # Use the first linked material as target
                target_material = linked_mats[0]
                
                # Remap all appended materials to the linked one
                for appended_mat in appended_mats:
                    if appended_mat.users > 0:
                        print(f"Remapping appended {appended_mat.name} ({appended_mat.users} users) to linked {target_material.name}")
                        appended_mat.user_remap(target_material)
                        remapping_count += 1
                    
                    # Remove the unused appended material
                    if appended_mat.users == 0:
                        print(f"Removing unused appended material: {appended_mat.name}")
                        bpy.data.materials.remove(appended_mat)
    
    # Also check for any BSDF materials that might be from old paths or different files
    # and try to find matching linked materials
    for mat in materials:
        if mat.library is None and (mat.name.startswith("BSDF_") or "BSDF_" in mat.name):
            # This is an appended BSDF material - look for a linked version
            base_name = mat.name.split(".blend")[0] if ".blend" in mat.name else mat.name
            base_name = base_name.split(".")[0] if "." in base_name else base_name
            
            # Look for any linked material with the same base name
            for linked_mat in materials:
                if (linked_mat.library is not None and 
                    linked_mat.name.startswith("BSDF_") and
                    (linked_mat.name == base_name or 
                     linked_mat.name.startswith(base_name + ".") or
                     linked_mat.name == mat.name)):
                    
                    if mat.users > 0:
                        print(f"Remapping old BSDF {mat.name} ({mat.users} users) to linked {linked_mat.name}")
                        mat.user_remap(linked_mat)
                        remapping_count += 1
                    
                    # Remove the unused material
                    if mat.users == 0:
                        print(f"Removing unused old BSDF material: {mat.name}")
                        bpy.data.materials.remove(mat)
                    break
    
    print(f"Remapped {remapping_count} appended/old BSDF materials to linked versions")
    return remapping_count

def remap_missing_datablocks():
    """Remap materials that have missing/broken library links"""
    
    print("\nChecking for missing datablocks to remap...")
    
    materials = bpy.data.materials
    remapping_count = 0
    
    # Find materials with missing library links
    missing_materials = []
    for mat in materials:
        if mat.library is not None and mat.library.filepath and not os.path.exists(bpy.path.abspath(mat.library.filepath)):
            missing_materials.append(mat)
            print(f"Found missing datablock: {mat.name} (from {mat.library.filepath})")
    
    if not missing_materials:
        print("No missing datablocks found.")
        return 0
    
    # For each missing material, try to find a replacement
    for missing_mat in missing_materials:
        base_name = missing_mat.name.split(".blend")[0] if ".blend" in missing_mat.name else missing_mat.name
        base_name = base_name.split(".")[0] if "." in base_name else base_name
        
        # Look for a replacement material
        replacement_found = False
        
        # First, try to find a linked material with the same name from the current library
        for mat in materials:
            if (mat.library is not None and 
                mat.library.filepath and 
                os.path.exists(bpy.path.abspath(mat.library.filepath)) and
                mat.name == missing_mat.name):
                
                if missing_mat.users > 0:
                    print(f"Remapping missing {missing_mat.name} ({missing_mat.users} users) to valid linked {mat.name}")
                    missing_mat.user_remap(mat)
                    remapping_count += 1
                replacement_found = True
                break
        
        # If no exact match, try to find a BSDF material with similar name
        if not replacement_found and (missing_mat.name.startswith("BSDF_") or "BSDF_" in missing_mat.name):
            for mat in materials:
                if (mat.library is not None and 
                    mat.library.filepath and 
                    os.path.exists(bpy.path.abspath(mat.library.filepath)) and
                    mat.name.startswith("BSDF_") and
                    (mat.name == base_name or 
                     mat.name.startswith(base_name + ".") or
                     base_name in mat.name)):
                    
                    if missing_mat.users > 0:
                        print(f"Remapping missing BSDF {missing_mat.name} ({missing_mat.users} users) to valid linked {mat.name}")
                        missing_mat.user_remap(mat)
                        remapping_count += 1
                    replacement_found = True
                    break
        
        # If still no replacement, try to find any valid linked material with the same base name
        if not replacement_found:
            for mat in materials:
                if (mat.library is not None and 
                    mat.library.filepath and 
                    os.path.exists(bpy.path.abspath(mat.library.filepath)) and
                    mat.name == base_name):
                    
                    if missing_mat.users > 0:
                        print(f"Remapping missing {missing_mat.name} ({missing_mat.users} users) to valid linked {mat.name}")
                        missing_mat.user_remap(mat)
                        remapping_count += 1
                    replacement_found = True
                    break
        
        if not replacement_found:
            print(f"Warning: No replacement found for missing material {missing_mat.name}")
    
    print(f"Remapped {remapping_count} missing datablocks to valid linked materials")
    return remapping_count

def replace_cel_materials():
    """Replace all CEL materials with their BSDF counterparts using Blender's user remapping"""
    
    # First, link BSDF materials from library
    linked_materials = link_bsdf_materials()
    
    # Then, remap any missing datablocks
    missing_remaps = remap_missing_datablocks()
    
    # Then, remap any appended BSDF materials to linked versions
    appended_remaps = remap_appended_to_linked()
    
    print(f"\n=== STARTING MATERIAL REPLACEMENT ===")
    
    # Custom material mappings (source -> target)
    custom_mappings = {
        "bag BLACK (squid ink)": "BSDF_black_SQUID-INK",
        "bag WHITE": "BSDF_WHITE",
        "Wheel-White": "BSDF_WHITE",
        "Bag Colors": "BSDF_Bag Colors",
        "cardboard": "Package_Cardboard",
        "blue (triton)": "BSDF_blue-2_TRITON",
        "gray (snow)": "BSDF_gray-6_SNOW",
        "gray (storm)": "BSDF_gray-2_STORM",
        "gray (summit)": "BSDF_gray-5_SUMMIT",
        "light blue (prime)": "BSDF_blue-4_PRIME",
        "yellow (summer)": "BSDF_orange-5_SUMMER",
        "Accessory_CEL_gray-6_SNOW": "BSDF_gray-6_SNOW",
        "Accessory_CEL_SquidInk": "BSDF_black_SQUID-INK",
        "FingerScanner": "BSDF_black_SQUID-INK",
        "cel BLACK (squid ink)": "BSDF_black_SQUID-INK",
        "cel WHITE": "BSDF_WHITE",
        "gray (stone)": "BSDF_gray-3_STONE",
        "green (oxygen)": "BSDF_green-3_OXYGEN",
        "orange (smile)": "BSDF_orange-3_SMILE"
    }
    
    # Get all materials in the scene
    materials = bpy.data.materials
    
    # Dictionary to store source -> target material mapping
    material_mapping = {}
    
    # Replace all CEL materials with their BSDF counterparts, ignoring numeric suffixes
    cel_pattern = re.compile(r"^(CEL_.+?)(\.\d{3})?$")
    bsdf_pattern = re.compile(r"^(BSDF_.+?)(\.\d{3})?$")
    
    # Build a mapping from base BSDF name to BSDF material (without suffix)
    bsdf_base_map = {bsdf_pattern.match(mat.name).group(1): mat for mat in materials if bsdf_pattern.match(mat.name)}
    
    # Build a mapping from exact material names to materials
    exact_material_map = {mat.name: mat for mat in materials}
    
    replacements_made = 0
    missing_targets = []
    
    # Process custom mappings first
    for source_name, target_name in custom_mappings.items():
        if source_name in exact_material_map:
            if target_name in exact_material_map:
                material_mapping[exact_material_map[source_name]] = exact_material_map[target_name]
                print(f"Found custom mapping: {source_name} -> {target_name}")
            else:
                missing_targets.append(f"{source_name} -> {target_name}")
                print(f"Warning: Target material '{target_name}' not found for custom mapping '{source_name}'")
    
    # Find all CEL materials and their BSDF counterparts
    for mat in materials:
        cel_match = cel_pattern.match(mat.name)
        if cel_match:
            base_cel = cel_match.group(1)
            base_bsdf = base_cel.replace("CEL_", "BSDF_", 1)
            if base_bsdf in bsdf_base_map:
                material_mapping[mat] = bsdf_base_map[base_bsdf]
                print(f"Found CEL mapping: {mat.name} -> {bsdf_base_map[base_bsdf].name}")
            else:
                missing_targets.append(f"{mat.name} -> {base_bsdf}")
                print(f"Warning: No BSDF counterpart found for {mat.name}")

    # Use Blender's built-in user remapping to replace ALL users
    for source_material, target_material in material_mapping.items():
        print(f"Remapping all users of {source_material.name} to {target_material.name}")
        
        # Get user count before remapping
        users_before = source_material.users
        
        # Remap all users of the source material to the target material
        # This catches ALL users including geometry node instances, drivers, etc.
        source_material.user_remap(target_material)
        
        # Get user count after remapping
        users_after = source_material.users
        
        replacements_this_material = users_before - users_after
        replacements_made += replacements_this_material
        
        print(f"  Users before: {users_before}, after: {users_after}")
        print(f"  Remapped {replacements_this_material} users")
    
    # Optional: Remove unused source materials after remapping
    print(f"\nCleaning up unused source materials...")
    removed_materials = 0
    
    for source_material in material_mapping.keys():
        if source_material.users == 0:
            print(f"Removing unused material: {source_material.name}")
            bpy.data.materials.remove(source_material)
            removed_materials += 1
        else:
            print(f"Warning: {source_material.name} still has {source_material.users} users after remapping")
    
    # Summary
    print(f"\n=== REPLACEMENT SUMMARY ===")
    print(f"Materials linked from library: {len(linked_materials)}")
    print(f"Missing datablocks remapped: {missing_remaps}")
    print(f"Appended->Linked remappings: {appended_remaps}")
    print(f"Total materials mapped: {len(material_mapping)}")
    print(f"Successful mappings: {len(material_mapping)}")
    print(f"Total user remappings: {replacements_made}")
    print(f"Source materials removed: {removed_materials}")
    
    if missing_targets:
        print(f"\nMissing target materials for:")
        for mapping in missing_targets:
            print(f"  - {mapping}")
    
    return len(material_mapping), replacements_made

# Run the replacement
if __name__ == "__main__":
    replace_cel_materials()

print("\nRemaining CEL materials in file:")
cel_count = 0
for mat in bpy.data.materials:
    if mat.name.startswith("CEL_"):
        print(f"  {mat.name} ({mat.users} users)")
        cel_count += 1

if cel_count == 0:
    print("  None - all CEL materials have been replaced!")

print("\nBSDF materials in file:")
for mat in bpy.data.materials:
    if mat.name.startswith("BSDF_"):
        print(f"  {mat.name} ({mat.users} users)")