#!/usr/bin/env python3
"""
optimization_engine Migration Script
=====================================
Automatically updates all imports across the codebase.

Usage:
    python migrate_imports.py --dry-run    # Preview changes
    python migrate_imports.py --execute    # Apply changes
"""

import os
import re
import sys
from pathlib import Path
from typing import Dict, List, Tuple

# Import mappings (old -> new) - using regex patterns
IMPORT_MAPPINGS = {
    # =============================================================================
    # CORE MODULE
    # =============================================================================
    r'from optimization_engine\.runner\b': 'from optimization_engine.core.runner',
    r'from optimization_engine\.base_runner\b': 'from optimization_engine.core.base_runner',
    r'from optimization_engine\.runner_with_neural\b': 'from optimization_engine.core.runner_with_neural',
    r'from optimization_engine\.intelligent_optimizer\b': 'from optimization_engine.core.intelligent_optimizer',
    r'from optimization_engine\.method_selector\b': 'from optimization_engine.core.method_selector',
    r'from optimization_engine\.strategy_selector\b': 'from optimization_engine.core.strategy_selector',
    r'from optimization_engine\.strategy_portfolio\b': 'from optimization_engine.core.strategy_portfolio',
    r'from optimization_engine\.gradient_optimizer\b': 'from optimization_engine.core.gradient_optimizer',
    r'import optimization_engine\.runner\b': 'import optimization_engine.core.runner',
    r'import optimization_engine\.intelligent_optimizer\b': 'import optimization_engine.core.intelligent_optimizer',

    # =============================================================================
    # SURROGATES MODULE
    # =============================================================================
    r'from optimization_engine\.neural_surrogate\b': 'from optimization_engine.processors.surrogates.neural_surrogate',
    r'from optimization_engine\.generic_surrogate\b': 'from optimization_engine.processors.surrogates.generic_surrogate',
    r'from optimization_engine\.adaptive_surrogate\b': 'from optimization_engine.processors.surrogates.adaptive_surrogate',
    r'from optimization_engine\.simple_mlp_surrogate\b': 'from optimization_engine.processors.surrogates.simple_mlp_surrogate',
    r'from optimization_engine\.active_learning_surrogate\b': 'from optimization_engine.processors.surrogates.active_learning_surrogate',
    r'from optimization_engine\.surrogate_tuner\b': 'from optimization_engine.processors.surrogates.surrogate_tuner',
    r'from optimization_engine\.auto_trainer\b': 'from optimization_engine.processors.surrogates.auto_trainer',
    r'from optimization_engine\.training_data_exporter\b': 'from optimization_engine.processors.surrogates.training_data_exporter',

    # =============================================================================
    # NX MODULE
    # =============================================================================
    r'from optimization_engine\.nx_solver\b': 'from optimization_engine.nx.solver',
    r'from optimization_engine\.nx_updater\b': 'from optimization_engine.nx.updater',
    r'from optimization_engine\.nx_session_manager\b': 'from optimization_engine.nx.session_manager',
    r'from optimization_engine\.solve_simulation\b': 'from optimization_engine.nx.solve_simulation',
    r'from optimization_engine\.solve_simulation_simple\b': 'from optimization_engine.nx.solve_simulation_simple',
    r'from optimization_engine\.model_cleanup\b': 'from optimization_engine.nx.model_cleanup',
    r'from optimization_engine\.export_expressions\b': 'from optimization_engine.nx.export_expressions',
    r'from optimization_engine\.import_expressions\b': 'from optimization_engine.nx.import_expressions',
    r'from optimization_engine\.mesh_converter\b': 'from optimization_engine.nx.mesh_converter',
    r'import optimization_engine\.nx_solver\b': 'import optimization_engine.nx.solver',
    r'import optimization_engine\.nx_updater\b': 'import optimization_engine.nx.updater',

    # =============================================================================
    # STUDY MODULE
    # =============================================================================
    r'from optimization_engine\.study_creator\b': 'from optimization_engine.study.creator',
    r'from optimization_engine\.study_wizard\b': 'from optimization_engine.study.wizard',
    r'from optimization_engine\.study_state\b': 'from optimization_engine.study.state',
    r'from optimization_engine\.study_reset\b': 'from optimization_engine.study.reset',
    r'from optimization_engine\.study_continuation\b': 'from optimization_engine.study.continuation',
    r'from optimization_engine\.benchmarking_substudy\b': 'from optimization_engine.study.benchmarking',
    r'from optimization_engine\.generate_history_from_trials\b': 'from optimization_engine.study.history_generator',

    # =============================================================================
    # REPORTING MODULE
    # =============================================================================
    r'from optimization_engine\.generate_report\b': 'from optimization_engine.reporting.report_generator',
    r'from optimization_engine\.generate_report_markdown\b': 'from optimization_engine.reporting.markdown_report',
    r'from optimization_engine\.comprehensive_results_analyzer\b': 'from optimization_engine.reporting.results_analyzer',
    r'from optimization_engine\.visualizer\b': 'from optimization_engine.reporting.visualizer',
    r'from optimization_engine\.landscape_analyzer\b': 'from optimization_engine.reporting.landscape_analyzer',

    # =============================================================================
    # CONFIG MODULE
    # =============================================================================
    r'from optimization_engine\.config_manager\b': 'from optimization_engine.config.manager',
    r'from optimization_engine\.optimization_config_builder\b': 'from optimization_engine.config.builder',
    r'from optimization_engine\.optimization_setup_wizard\b': 'from optimization_engine.config.setup_wizard',
    r'from optimization_engine\.capability_matcher\b': 'from optimization_engine.config.capability_matcher',
    r'from optimization_engine\.template_loader\b': 'from optimization_engine.config.template_loader',

    # =============================================================================
    # UTILS MODULE
    # =============================================================================
    r'from optimization_engine\.logger\b': 'from optimization_engine.utils.logger',
    r'from optimization_engine\.auto_doc\b': 'from optimization_engine.utils.auto_doc',
    r'from optimization_engine\.realtime_tracking\b': 'from optimization_engine.utils.realtime_tracking',
    r'from optimization_engine\.codebase_analyzer\b': 'from optimization_engine.utils.codebase_analyzer',
    r'from optimization_engine\.pruning_logger\b': 'from optimization_engine.utils.pruning_logger',

    # =============================================================================
    # FUTURE MODULE
    # =============================================================================
    r'from optimization_engine\.research_agent\b': 'from optimization_engine.future.research_agent',
    r'from optimization_engine\.pynastran_research_agent\b': 'from optimization_engine.future.pynastran_research_agent',
    r'from optimization_engine\.targeted_research_planner\b': 'from optimization_engine.future.targeted_research_planner',
    r'from optimization_engine\.workflow_decomposer\b': 'from optimization_engine.future.workflow_decomposer',
    r'from optimization_engine\.step_classifier\b': 'from optimization_engine.future.step_classifier',
    r'from optimization_engine\.llm_optimization_runner\b': 'from optimization_engine.future.llm_optimization_runner',
    r'from optimization_engine\.llm_workflow_analyzer\b': 'from optimization_engine.future.llm_workflow_analyzer',

    # =============================================================================
    # EXTRACTORS/VALIDATORS additions
    # =============================================================================
    r'from optimization_engine\.op2_extractor\b': 'from optimization_engine.extractors.op2_extractor',
    r'from optimization_engine\.extractor_library\b': 'from optimization_engine.extractors.extractor_library',
    r'from optimization_engine\.simulation_validator\b': 'from optimization_engine.validators.simulation_validator',

    # =============================================================================
    # PROCESSORS
    # =============================================================================
    r'from optimization_engine\.adaptive_characterization\b': 'from optimization_engine.processors.adaptive_characterization',
}

# Also need to handle utils submodule imports that moved
UTILS_MAPPINGS = {
    r'from optimization_engine\.utils\.nx_session_manager\b': 'from optimization_engine.nx.session_manager',
}

# Combine all mappings
ALL_MAPPINGS = {**IMPORT_MAPPINGS, **UTILS_MAPPINGS}

def find_files(root: Path, extensions: List[str], exclude_dirs: List[str] = None) -> List[Path]:
    """Find all files with given extensions, excluding certain directories."""
    if exclude_dirs is None:
        exclude_dirs = ['optimization_engine_BACKUP', '.venv', 'node_modules', '__pycache__', '.git']

    files = []
    for ext in extensions:
        for f in root.rglob(f'*{ext}'):
            # Check if any excluded dir is in the path
            if not any(excl in str(f) for excl in exclude_dirs):
                files.append(f)
    return files

def update_file(filepath: Path, mappings: Dict[str, str], dry_run: bool = True) -> Tuple[int, List[str]]:
    """Update imports in a single file."""
    try:
        content = filepath.read_text(encoding='utf-8', errors='ignore')
    except Exception as e:
        print(f"  ERROR reading {filepath}: {e}")
        return 0, []

    changes = []
    new_content = content

    for pattern, replacement in mappings.items():
        matches = re.findall(pattern, content)
        if matches:
            new_content = re.sub(pattern, replacement, new_content)
            changes.append(f"  {pattern} -> {replacement} ({len(matches)} occurrences)")

    if changes and not dry_run:
        filepath.write_text(new_content, encoding='utf-8')

    return len(changes), changes

def main():
    dry_run = '--dry-run' in sys.argv or '--execute' not in sys.argv

    if dry_run:
        print("=" * 60)
        print("DRY RUN MODE - No files will be modified")
        print("=" * 60)
    else:
        print("=" * 60)
        print("EXECUTE MODE - Files will be modified!")
        print("=" * 60)
        confirm = input("Are you sure? (yes/no): ")
        if confirm.lower() != 'yes':
            print("Aborted.")
            return

    root = Path('.')

    # Find all Python files
    py_files = find_files(root, ['.py'])
    print(f"\nFound {len(py_files)} Python files to check")

    total_changes = 0
    files_changed = 0

    for filepath in sorted(py_files):
        count, changes = update_file(filepath, ALL_MAPPINGS, dry_run)
        if count > 0:
            files_changed += 1
            total_changes += count
            print(f"\n{filepath} ({count} changes):")
            for change in changes:
                print(change)

    print("\n" + "=" * 60)
    print(f"SUMMARY: {total_changes} changes in {files_changed} files")
    print("=" * 60)

    if dry_run:
        print("\nTo apply changes, run: python migrate_imports.py --execute")

if __name__ == '__main__':
    main()