"""
Automatic Levy Flight Integration Script
========================================

This script automatically adds Levy flight support to all existing algorithms
in the algorithms/ directory without manually modifying each file.

Usage:
    python -m mha_toolbox.apply_levy_to_all

Author: MHA Flow Team
License: MIT
"""

import os
import ast
import sys
from pathlib import Path
from typing import List, Tuple


class LevyFlightIntegrator:
    """Automatically integrate Levy flight into algorithm files"""
    
    def __init__(self, algorithms_dir: str):
        self.algorithms_dir = Path(algorithms_dir)
        self.modified_files = []
        self.skipped_files = []
        self.errors = []
    
    def process_all_algorithms(self) -> dict:
        """
        Process all algorithm files in the algorithms directory.
        
        Returns
        -------
        dict
            Statistics about the integration process
        """
        print("=" * 70)
        print("LEVY FLIGHT UNIVERSAL INTEGRATION")
        print("=" * 70)
        print(f"\nScanning directory: {self.algorithms_dir}")
        
        # Find all Python files
        py_files = list(self.algorithms_dir.rglob("*.py"))
        py_files = [f for f in py_files if f.name != "__init__.py"]
        
        print(f"Found {len(py_files)} algorithm files")
        print("\nProcessing...\n")
        
        for file_path in py_files:
            try:
                self._process_file(file_path)
            except Exception as e:
                self.errors.append((file_path, str(e)))
                print(f"  ❌ ERROR: {file_path.name} - {str(e)}")
        
        return self._generate_report()
    
    def _process_file(self, file_path: Path):
        """Process a single algorithm file"""
        # Read file content
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # Check if already has Levy flight import
        if 'levy_flight_universal' in content or 'LevyFlightMixin' in content:
            self.skipped_files.append(file_path)
            print(f"  ⏭️  SKIP: {file_path.name} (already has Levy flight)")
            return
        
        # Parse AST to check if it's an optimizer class
        try:
            tree = ast.parse(content)
        except SyntaxError:
            self.skipped_files.append(file_path)
            print(f"  ⏭️  SKIP: {file_path.name} (syntax error)")
            return
        
        # Find optimizer classes
        optimizer_classes = self._find_optimizer_classes(tree)
        
        if not optimizer_classes:
            self.skipped_files.append(file_path)
            print(f"  ⏭️  SKIP: {file_path.name} (no optimizer class found)")
            return
        
        # Add import statement
        modified_content = self._add_levy_import(content)
        
        # Write back to file
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(modified_content)
        
        self.modified_files.append(file_path)
        print(f"  ✅ MODIFIED: {file_path.name}")
    
    def _find_optimizer_classes(self, tree: ast.AST) -> List[str]:
        """Find all optimizer class names in the AST"""
        classes = []
        
        for node in ast.walk(tree):
            if isinstance(node, ast.ClassDef):
                # Check if it inherits from BaseOptimizer or similar
                for base in node.bases:
                    if isinstance(base, ast.Name):
                        if 'Optimizer' in base.id or 'Algorithm' in base.id:
                            classes.append(node.name)
                            break
        
        return classes
    
    def _add_levy_import(self, content: str) -> str:
        """
        Add Levy flight import to the file content.
        
        Adds the import after existing imports.
        """
        lines = content.split('\n')
        
        # Find the last import line
        last_import_idx = 0
        for i, line in enumerate(lines):
            stripped = line.strip()
            if stripped.startswith('import ') or stripped.startswith('from '):
                last_import_idx = i
        
        # Add the Levy flight import
        levy_import = "from .levy_flight_universal import add_levy_flight_to_position"
        
        # Check if we need relative import based on file location
        if '/hybrid/' in content or '\\hybrid\\' in content:
            levy_import = "from ..levy_flight_universal import add_levy_flight_to_position"
        
        # Insert the import
        lines.insert(last_import_idx + 1, levy_import)
        
        # Add usage comment in docstring or after class definition
        modified_lines = []
        added_comment = False
        
        for i, line in enumerate(lines):
            modified_lines.append(line)
            
            # Add comment after _optimize method definition
            if not added_comment and 'def _optimize(' in line:
                indent = len(line) - len(line.lstrip())
                comment = ' ' * (indent + 4) + '# Levy flight available via: add_levy_flight_to_position()'
                modified_lines.append(comment)
                added_comment = True
        
        return '\n'.join(modified_lines)
    
    def _generate_report(self) -> dict:
        """Generate integration report"""
        print("\n" + "=" * 70)
        print("INTEGRATION REPORT")
        print("=" * 70)
        
        report = {
            'total_files': len(self.modified_files) + len(self.skipped_files) + len(self.errors),
            'modified': len(self.modified_files),
            'skipped': len(self.skipped_files),
            'errors': len(self.errors)
        }
        
        print(f"\n📊 Statistics:")
        print(f"  Total Files Scanned: {report['total_files']}")
        print(f"  ✅ Modified: {report['modified']}")
        print(f"  ⏭️  Skipped: {report['skipped']}")
        print(f"  ❌ Errors: {report['errors']}")
        
        if self.modified_files:
            print(f"\n✅ Modified Files ({len(self.modified_files)}):")
            for file_path in self.modified_files:
                print(f"  - {file_path.relative_to(self.algorithms_dir)}")
        
        if self.errors:
            print(f"\n❌ Errors ({len(self.errors)}):")
            for file_path, error in self.errors:
                print(f"  - {file_path.name}: {error}")
        
        print("\n" + "=" * 70)
        print("INTEGRATION COMPLETE!")
        print("=" * 70)
        print("\n💡 Usage in algorithms:")
        print("```python")
        print("# In your _optimize method:")
        print("if iteration % 5 == 0:")
        print("    new_position = add_levy_flight_to_position(")
        print("        position=individual,")
        print("        best_position=best_solution,")
        print("        bounds=bounds,")
        print("        iteration=iteration,")
        print("        max_iterations=self.max_iterations")
        print("    )")
        print("```")
        
        return report


def main():
    """Main execution function"""
    # Get the algorithms directory
    script_dir = Path(__file__).parent
    algorithms_dir = script_dir / "algorithms"
    
    if not algorithms_dir.exists():
        print(f"❌ Error: Algorithms directory not found at {algorithms_dir}")
        sys.exit(1)
    
    # Create integrator and process all files
    integrator = LevyFlightIntegrator(algorithms_dir)
    report = integrator.process_all_algorithms()
    
    # Return status code
    if report['errors'] > 0:
        sys.exit(1)
    else:
        sys.exit(0)


if __name__ == "__main__":
    main()
