#!/usr/bin/env python3
"""
Fix missing slot imports in class files - Version 2.

This script properly handles YAML indentation when adding imports.
"""

import re
from pathlib import Path
import yaml

SCHEMA_DIR = Path('/Users/kempersc/apps/glam/schemas/20251121/linkml')
CLASSES_DIR = SCHEMA_DIR / 'modules' / 'classes'
SLOTS_DIR = SCHEMA_DIR / 'modules' / 'slots'

# Get all available slot files
available_slots = {f.stem for f in SLOTS_DIR.glob('*.yaml')}
print(f"Available slot files: {len(available_slots)}")

def get_class_slots(content: str) -> set:
    """Extract slot names from a class file's slots: list."""
    slots = set()
    try:
        data = yaml.safe_load(content)
        if data and 'classes' in data:
            for class_name, class_def in data['classes'].items():
                if class_def and 'slots' in class_def:
                    for slot in class_def['slots']:
                        slots.add(slot)
    except yaml.YAMLError:
        return set()
    return slots

def get_imported_slots(content: str) -> set:
    """Extract slot imports from a class file."""
    imported = set()
    # Match imports like: - ../slots/slot_name
    pattern = r'-\s*\.\./slots/(\w+)'
    for match in re.finditer(pattern, content):
        imported.add(match.group(1))
    return imported

def add_slot_imports(content: str, slots_to_add: set) -> str:
    """Add missing slot imports to content, preserving YAML structure."""
    lines = content.split('\n')
    result_lines = []
    
    in_imports = False
    imports_indent = ''
    last_import_line_idx = -1
    
    for i, line in enumerate(lines):
        result_lines.append(line)
        
        # Detect start of imports section
        if re.match(r'^imports:\s*$', line):
            in_imports = True
            continue
        
        # Inside imports section
        if in_imports:
            # Check if this is an import line (starts with - after whitespace)
            match = re.match(r'^(\s*)-\s*', line)
            if match:
                imports_indent = match.group(1)
                last_import_line_idx = len(result_lines)
            elif line.strip() and not line.strip().startswith('#'):
                # Non-empty, non-comment line that's not an import - we've left imports section
                in_imports = False
    
    # Insert new imports after the last import line
    # Note: imports_indent can be '' (empty string) which is falsy but valid
    if last_import_line_idx > 0:
        # Create import lines with correct indentation
        new_imports = [f"{imports_indent}- ../slots/{slot}" for slot in sorted(slots_to_add)]
        # Insert after last import
        for new_import in reversed(new_imports):
            result_lines.insert(last_import_line_idx, new_import)
    
    return '\n'.join(result_lines)

def main():
    total_fixed = 0
    files_modified = 0
    errors = 0
    
    for class_file in sorted(CLASSES_DIR.glob('*.yaml')):
        content = class_file.read_text()
        
        # Get slots used by this class
        used_slots = get_class_slots(content)
        if not used_slots:
            continue
        
        # Get already imported slots
        imported_slots = get_imported_slots(content)
        
        # Find slots that need to be imported
        missing_imports = used_slots - imported_slots
        
        # Filter to only slots that have files
        valid_missing = missing_imports & available_slots
        
        if valid_missing:
            new_content = add_slot_imports(content, valid_missing)
            
            # Verify the new content is valid YAML
            try:
                yaml.safe_load(new_content)
                class_file.write_text(new_content)
                print(f"OK: {class_file.name}: Added {len(valid_missing)} imports")
                total_fixed += len(valid_missing)
                files_modified += 1
            except yaml.YAMLError as e:
                print(f"ERROR: {class_file.name}: YAML error after modification: {e}")
                errors += 1
    
    print(f"\n=== Summary ===")
    print(f"Files modified: {files_modified}")
    print(f"Total imports added: {total_fixed}")
    print(f"Errors: {errors}")

if __name__ == '__main__':
    main()