import os
import re
from safetensors.torch import load_file, save_file

def convert_key(key):
    # 1. Remove the original prefix
    new_key = key.replace("lora_unet__", "")
    
    # 2. Identify the parts.
    parts = new_key.split(".")
    block_name = parts[0]  # e.g., "layers_0_attention_to_k"
    rest = parts[1:]       # e.g., ["lora_down", "weight"]

    # --- FIX 1: Layers and Refiners (Force dot after number) ---
    # The previous script missed the trailing underscore. 
    # We replace "layers_0_" with "layers.0."
    
    # Fix "layers_X_" -> "layers.X."
    block_name = re.sub(r'layers_(\d+)_', r'layers.\1.', block_name)
    
    # Fix "context_refiner_X_" -> "context_refiner.X."
    block_name = re.sub(r'context_refiner_(\d+)_', r'context_refiner.\1.', block_name)
    
    # Fix "noise_refiner_X_" -> "noise_refiner.X."
    block_name = re.sub(r'noise_refiner_(\d+)_', r'noise_refiner.\1.', block_name)

    # --- FIX 2: Attention blocks ---
    
    # Fix "attention_to_k" -> "attention.to_k"
    for t in ["to_k", "to_q", "to_v"]:
        if f"_{t}" in block_name:
            block_name = block_name.replace(f"_{t}", f".{t}")
            
    # Fix "attention_to_out_0" -> "attention.to_out.0"
    # This is specific: ComfyUI expects "to_out.0"
    if "_to_out" in block_name:
        block_name = block_name.replace("_to_out", ".to_out")
        # Handle the trailing number for output (e.g. to_out_0 -> to_out.0)
        block_name = re.sub(r'\.to_out_(\d+)', r'.to_out.\1', block_name)

    # --- FINALIZE ---
    final_key = "diffusion_model." + ".".join([block_name] + rest)
    return final_key

# --- MAIN EXECUTION ---
print("Looking for .safetensors files...")
# Find original file (exclude previously converted ones to be safe)
files = [f for f in os.listdir('.') if f.endswith('.safetensors') and 'converted' not in f and 'fixed' not in f]

if not files:
    print("Error: No original .safetensors file found.")
    print("Make sure the original 'Z-Image-Fun-Lora-Distill-8-Steps.safetensors' is in this folder.")
    exit()

input_file = files[0]
print(f"Processing: {input_file}")

try:
    tensors = load_file(input_file)
    new_tensors = {}

    print("Converting keys...")
    # Print a preview of the first converted key to verify
    first_key = list(tensors.keys())[0]
    print(f"Preview: {first_key} \n      -> {convert_key(first_key)}")

    for k, v in tensors.items():
        new_k = convert_key(k)
        new_tensors[new_k] = v
        
    output_file = input_file.replace(".safetensors", "_v3_fixed.safetensors")
    save_file(new_tensors, output_file)

    print(f"\nSUCCESS! Created: {output_file}")
    print("Move this file to ComfyUI/models/loras/ and DELETE the old converted ones.")

except Exception as e:
    print(f"An error occurred: {e}")