import os

def reorder_imports(path):
    with open(path, 'r') as f:
        lines = f.readlines()
        
    # Find imports block
    start_idx = -1
    end_idx = -1
    
    for i, line in enumerate(lines):
        if line.strip() == "imports:":
            start_idx = i
            break
            
    if start_idx == -1:
        return

    for i in range(start_idx + 1, len(lines)):
        line = lines[i]
        if line.strip() and not line.startswith(" ") and not line.startswith("#"):
             end_idx = i
             break
    
    if end_idx == -1:
        end_idx = len(lines)
        
    imports = lines[start_idx + 1:end_idx]
    
    # Separate class imports
    class_imports = []
    other_imports = []
    
    for line in imports:
        if "modules/classes/" in line:
            class_imports.append(line)
        else:
            other_imports.append(line)
            
    # Define priority classes
    priority_classes = [
        "ReconstructedEntity",
        "CustodianType",
        "Activity",
        "Event",
        "Provenance",
        "Identifier",
        "Label",
        "Description",
        "Content",
        "CollectionType",
        "FindingAidType",
        "DocumentType",
        "StatementType",
        "NameType",
        "LabelType",
        "Place",
        "TimeSpan",
        "Agent",
        "Organization",
        "Person",
        "Group",
        "Thing",
        "Concept",
        "InformationObject",
        "DigitalPlatform",
        "Storage",
        "FinancialStatement",
        "Budget",
        "Plan",
        "Action"
    ]
    
    sorted_class_imports = []
    remaining_class_imports = []
    
    # Extract priority
    for p in priority_classes:
        import_line = f"- modules/classes/{p}\n"
        # Find matching line (ignoring whitespace)
        found = False
        # Iterate over copy to allow removal
        for line in list(class_imports):
            if p in line and line.strip().endswith(p): # Endswith classname
                 sorted_class_imports.append(line)
                 class_imports.remove(line)
                 found = True
                 break
        if not found:
            # Try to find without exact match
             pass
             
    # Sort remaining
    class_imports.sort()
    
    final_imports = other_imports + sorted_class_imports + class_imports
    
    new_lines = lines[:start_idx + 1] + final_imports + lines[end_idx:]
    
    with open(path, 'w') as f:
        f.writelines(new_lines)
        
    print(f"Reordered imports in {path}")

reorder_imports("schemas/20251121/linkml/01_custodian_name_modular.yaml")