import os

def populate_imports(main_schema_path, classes_dir):
    with open(main_schema_path, 'r') as f:
        lines = f.readlines()
        
    # Find start and end of imports section
    start_idx = -1
    end_idx = -1
    
    for i, line in enumerate(lines):
        if line.strip() == "imports:":
            start_idx = i
            break
            
    if start_idx == -1:
        print("imports: section not found")
        return

    # Find where imports block ends (next top-level key or end of file)
    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)
        
    # Keep existing non-class imports (slots, enums, metadata, types)
    # Filter out modules/classes/* imports
    
    preserved_imports = []
    for i in range(start_idx + 1, end_idx):
        line = lines[i]
        stripped = line.strip()
        if stripped.startswith("- modules/classes/"):
            continue # We will regenerate these
        if stripped.startswith("- ./"): # Should not happen in 01_...
             continue
        if stripped.strip():
             preserved_imports.append(line)
             
    # Collect all class files
    class_imports = []
    for root, dirs, files in os.walk(classes_dir):
        if 'archive' in dirs:
            dirs.remove('archive')
        if 'deprecated' in dirs:
            dirs.remove('deprecated')
        
        if 'archive' in root.split(os.sep): 
             continue
        if 'deprecated' in root.split(os.sep): 
             continue
             
        for file in files:
            if file.endswith(".yaml"):
                # Path relative to 01_... location (schemas/20251121/linkml)
                # classes_dir is schemas/20251121/linkml/modules/classes
                # relative path: modules/classes/Filename
                
                # Check for duplicate base names? e.g. .yaml vs .yaml.bak?
                # assume clean dir.
                
                rel_path = f"modules/classes/{file.replace('.yaml', '')}"
                class_imports.append(f"  - {rel_path}\n")
    
    class_imports.sort()
    
    # Reassemble file
    new_lines = lines[:start_idx + 1] + preserved_imports + class_imports + lines[end_idx:]
    
    with open(main_schema_path, 'w') as f:
        f.writelines(new_lines)
        
    print(f"Updated {main_schema_path} with {len(class_imports)} class imports.")

populate_imports("schemas/20251121/linkml/01_custodian_name_modular.yaml", "schemas/20251121/linkml/modules/classes")