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

This script scans class files for slots used in the `slots:` list
and adds missing imports for those slots.
"""

import os
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(class_file: Path) -> set:
    """Extract slot names from a class file's slots: list."""
    with open(class_file) as f:
        content = f.read()
    
    # Find slots used in classes
    slots = set()
    
    # Pattern to find slots: section under classes
    # This is tricky because we need to parse YAML structure
    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 as e:
        print(f"YAML error in {class_file}: {e}")
        return set()
    
    return slots

def get_imported_slots(class_file: Path) -> set:
    """Extract slot imports from a class file."""
    with open(class_file) as f:
        content = f.read()
    
    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(class_file: Path, slots_to_add: set):
    """Add missing slot imports to a class file."""
    with open(class_file) as f:
        content = f.read()
    
    # Find the imports section
    lines = content.split('\n')
    new_lines = []
    in_imports = False
    imports_ended = False
    added = False
    
    for i, line in enumerate(lines):
        new_lines.append(line)
        
        # Detect imports section
        if line.strip() == 'imports:':
            in_imports = True
            continue
        
        if in_imports and not imports_ended:
            # Check if we're still in imports (lines starting with -)
            if line.strip().startswith('-'):
                continue
            else:
                # End of imports section - add our slots before this line
                imports_ended = True
                # Insert slot imports before the current line
                for slot in sorted(slots_to_add):
                    new_lines.insert(len(new_lines) - 1, f'- ../slots/{slot}')
                added = True
    
    # If we didn't find a good place, add after last import
    if not added:
        # Find imports section and add at the end
        result_lines = []
        in_imports = False
        last_import_idx = -1
        
        for i, line in enumerate(lines):
            result_lines.append(line)
            if line.strip() == 'imports:':
                in_imports = True
            elif in_imports and line.strip().startswith('-'):
                last_import_idx = len(result_lines)
            elif in_imports and not line.strip().startswith('-') and line.strip() != '':
                in_imports = False
        
        if last_import_idx > 0:
            # Insert after last import
            for slot in sorted(slots_to_add, reverse=True):
                result_lines.insert(last_import_idx, f'- ../slots/{slot}')
            new_lines = result_lines
    
    with open(class_file, 'w') as f:
        f.write('\n'.join(new_lines))

def main():
    total_fixed = 0
    files_modified = 0
    
    for class_file in sorted(CLASSES_DIR.glob('*.yaml')):
        # Get slots used by this class
        used_slots = get_class_slots(class_file)
        if not used_slots:
            continue
        
        # Get already imported slots
        imported_slots = get_imported_slots(class_file)
        
        # 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:
            print(f"{class_file.name}: Adding {len(valid_missing)} slot imports")
            add_slot_imports(class_file, valid_missing)
            total_fixed += len(valid_missing)
            files_modified += 1
    
    print(f"\n=== Summary ===")
    print(f"Files modified: {files_modified}")
    print(f"Total imports added: {total_fixed}")

if __name__ == '__main__':
    main()