"""
Custom Algorithm Manager
========================

This module handles uploading, validating, and integrating custom MHA algorithms.

Features:
- Parse uploaded Python files
- Identify MHA algorithms automatically
- Extract parameters and map them
- Validate algorithm structure
- Integrate into the system
- Support for GitHub contributions

Author: MHA Flow Development Team
License: MIT
"""

import ast
import inspect
import importlib.util
import os
import json
import re
from typing import Dict, List, Tuple, Optional, Any
from pathlib import Path
import numpy as np

# Import GitHub integration
try:
    from mha_toolbox.github_integration import GitHubAlgorithmUploader
    GITHUB_AVAILABLE = True
except ImportError:
    GITHUB_AVAILABLE = False


class AlgorithmParser:
    """Parse and analyze uploaded Python files to extract algorithm information"""
    
    def __init__(self):
        self.required_methods = ['optimize', '__init__']
        self.optional_methods = ['update', 'evaluate', 'initialize_population']
    
    def parse_file(self, file_path: str) -> Dict[str, Any]:
        """
        Parse a Python file and extract algorithm information.
        
        Parameters
        ----------
        file_path : str
            Path to the Python file
            
        Returns
        -------
        info : dict
            Dictionary containing:
            - classes: List of class definitions
            - functions: List of function definitions
            - imports: List of imports
            - is_valid_mha: Whether it's a valid MHA algorithm
        """
        with open(file_path, 'r', encoding='utf-8') as f:
            source_code = f.read()
        
        try:
            tree = ast.parse(source_code)
        except SyntaxError as e:
            return {
                'error': f"Syntax error in file: {e}",
                'is_valid_mha': False
            }
        
        classes = []
        functions = []
        imports = []
        
        for node in ast.walk(tree):
            if isinstance(node, ast.ClassDef):
                class_info = self._analyze_class(node)
                classes.append(class_info)
            elif isinstance(node, ast.FunctionDef):
                func_info = self._analyze_function(node)
                functions.append(func_info)
            elif isinstance(node, (ast.Import, ast.ImportFrom)):
                imports.append(self._get_import_info(node))
        
        # Determine if it's a valid MHA
        is_valid_mha, main_class = self._is_valid_mha(classes)
        
        return {
            'classes': classes,
            'functions': functions,
            'imports': imports,
            'is_valid_mha': is_valid_mha,
            'main_class': main_class,
            'source_code': source_code
        }
    
    def _analyze_class(self, node: ast.ClassDef) -> Dict:
        """Analyze a class definition"""
        methods = []
        for item in node.body:
            if isinstance(item, ast.FunctionDef):
                method_info = self._analyze_function(item)
                methods.append(method_info)
        
        # Get base classes
        bases = [self._get_name(base) for base in node.bases]
        
        return {
            'name': node.name,
            'methods': methods,
            'bases': bases,
            'lineno': node.lineno
        }
    
    def _analyze_function(self, node: ast.FunctionDef) -> Dict:
        """Analyze a function definition"""
        # Extract parameters
        params = []
        for arg in node.args.args:
            param_name = arg.arg
            # Try to get type annotation if available
            param_type = None
            if arg.annotation:
                param_type = ast.unparse(arg.annotation) if hasattr(ast, 'unparse') else None
            params.append({
                'name': param_name,
                'type': param_type
            })
        
        # Get return annotation
        return_type = None
        if node.returns:
            return_type = ast.unparse(node.returns) if hasattr(ast, 'unparse') else None
        
        # Extract docstring
        docstring = ast.get_docstring(node)
        
        return {
            'name': node.name,
            'params': params,
            'return_type': return_type,
            'docstring': docstring,
            'lineno': node.lineno
        }
    
    def _get_import_info(self, node) -> str:
        """Get import information"""
        if isinstance(node, ast.Import):
            return ', '.join([alias.name for alias in node.names])
        elif isinstance(node, ast.ImportFrom):
            module = node.module or ''
            names = ', '.join([alias.name for alias in node.names])
            return f"from {module} import {names}"
        return ''
    
    def _get_name(self, node) -> str:
        """Get name from AST node"""
        if isinstance(node, ast.Name):
            return node.id
        elif isinstance(node, ast.Attribute):
            return f"{self._get_name(node.value)}.{node.attr}"
        return str(node)
    
    def _is_valid_mha(self, classes: List[Dict]) -> Tuple[bool, Optional[Dict]]:
        """
        Determine if the file contains a valid MHA algorithm.
        
        A valid MHA must:
        1. Have at least one class
        2. Have an 'optimize' method
        3. Have an '__init__' method with parameters
        """
        for cls in classes:
            method_names = [m['name'] for m in cls['methods']]
            
            # Check for required methods
            has_optimize = 'optimize' in method_names
            has_init = '__init__' in method_names
            
            if has_optimize and has_init:
                # Get optimize method details
                optimize_method = next(m for m in cls['methods'] if m['name'] == 'optimize')
                
                # Check if optimize has expected parameters
                param_names = [p['name'] for p in optimize_method['params']]
                
                # Common MHA optimize parameters
                mha_indicators = ['objective_func', 'fitness_func', 'bounds', 
                                'dimensions', 'problem', 'n_dims']
                
                if any(indicator in param_names for indicator in mha_indicators):
                    return True, cls
        
        return False, None


class ParameterMapper:
    """Map custom algorithm parameters to standard system parameters"""
    
    def __init__(self):
        self.standard_params = {
            'population_size': ['pop_size', 'n_agents', 'n_particles', 'swarm_size', 
                              'colony_size', 'n_individuals', 'NOP'],
            'max_iterations': ['max_iter', 'iterations', 'n_iterations', 'max_gen',
                             'generations', 'epochs', 'MAX_ITER'],
            'dimensions': ['dim', 'n_dims', 'n_dimensions', 'D', 'problem_size'],
            'bounds': ['lb_ub', 'search_space', 'limits', 'min_max'],
            'objective_function': ['fitness_func', 'cost_func', 'eval_func', 
                                 'objective', 'fitness_function']
        }
    
    def suggest_mapping(self, param_name: str) -> Optional[str]:
        """
        Suggest a standard parameter name for a custom parameter.
        
        Parameters
        ----------
        param_name : str
            Custom parameter name
            
        Returns
        -------
        standard_name : str or None
            Suggested standard parameter name
        """
        param_lower = param_name.lower()
        
        for standard, variants in self.standard_params.items():
            if param_lower in [v.lower() for v in variants]:
                return standard
            # Check partial matches
            if any(variant.lower() in param_lower for variant in variants):
                return standard
        
        return None
    
    def create_mapping_prompt(self, detected_params: List[str]) -> Dict[str, Any]:
        """
        Create a mapping prompt for user to confirm parameter mappings.
        
        Parameters
        ----------
        detected_params : list
            List of detected parameter names
            
        Returns
        -------
        mapping_prompt : dict
            Dictionary with suggested mappings for UI display
        """
        mappings = {}
        unmapped = []
        
        for param in detected_params:
            suggestion = self.suggest_mapping(param)
            if suggestion:
                mappings[param] = {
                    'suggested': suggestion,
                    'confidence': 'high',
                    'confirmed': False
                }
            else:
                unmapped.append(param)
        
        return {
            'mapped': mappings,
            'unmapped': unmapped,
            'requires_user_input': len(unmapped) > 0
        }


class AlgorithmValidator:
    """Validate uploaded algorithms for correctness and compatibility"""
    
    def __init__(self):
        self.validation_rules = {
            'has_class': 'Algorithm must be defined as a class',
            'has_optimize': 'Must have an optimize() method',
            'has_init': 'Must have __init__() method',
            'returns_result': 'optimize() must return results',
            'handles_bounds': 'Must handle search space bounds',
            'numpy_compatible': 'Must work with numpy arrays'
        }
    
    def validate(self, algorithm_info: Dict) -> Dict[str, Any]:
        """
        Validate an algorithm's structure and compatibility.
        
        Parameters
        ----------
        algorithm_info : dict
            Parsed algorithm information
            
        Returns
        -------
        validation_result : dict
            Validation results with pass/fail for each rule
        """
        results = {}
        
        if not algorithm_info.get('is_valid_mha', False):
            return {
                'valid': False,
                'error': 'Not recognized as a valid MHA algorithm',
                'details': results
            }
        
        main_class = algorithm_info.get('main_class')
        if not main_class:
            return {
                'valid': False,
                'error': 'No main algorithm class found',
                'details': results
            }
        
        # Check each validation rule
        results['has_class'] = main_class is not None
        
        method_names = [m['name'] for m in main_class['methods']]
        results['has_optimize'] = 'optimize' in method_names
        results['has_init'] = '__init__' in method_names
        
        # Check optimize method signature
        if results['has_optimize']:
            optimize_method = next(m for m in main_class['methods'] if m['name'] == 'optimize')
            param_names = [p['name'] for p in optimize_method['params']]
            results['handles_bounds'] = any(b in param_names for b in ['bounds', 'lb_ub', 'limits'])
            
            # Check if returns something (has return statement in docstring or annotation)
            results['returns_result'] = (
                optimize_method.get('return_type') is not None or
                (optimize_method.get('docstring') and 'return' in optimize_method['docstring'].lower())
            )
        else:
            results['handles_bounds'] = False
            results['returns_result'] = False
        
        # Check for numpy usage
        imports = algorithm_info.get('imports', [])
        results['numpy_compatible'] = any('numpy' in imp or 'np' in imp for imp in imports)
        
        # Overall validation
        all_passed = all(results.values())
        
        return {
            'valid': all_passed,
            'details': results,
            'warnings': self._generate_warnings(results),
            'main_class_name': main_class['name']
        }
    
    def _generate_warnings(self, results: Dict[str, bool]) -> List[str]:
        """Generate warning messages for failed validations"""
        warnings = []
        for rule, passed in results.items():
            if not passed and rule in self.validation_rules:
                warnings.append(self.validation_rules[rule])
        return warnings


class AlgorithmIntegrator:
    """Integrate validated custom algorithms into the system"""
    
    def __init__(self, custom_algorithms_dir: str = "mha_toolbox/algorithms/custom"):
        self.custom_dir = Path(custom_algorithms_dir)
        self.custom_dir.mkdir(parents=True, exist_ok=True)
        self.registry_file = self.custom_dir / "custom_registry.json"
        
        # Initialize registry
        if not self.registry_file.exists():
            self._init_registry()
    
    def _init_registry(self):
        """Initialize empty registry"""
        registry = {
            'algorithms': {},
            'metadata': {
                'created': str(Path(__file__).stat().st_ctime),
                'version': '1.0.0'
            }
        }
        with open(self.registry_file, 'w') as f:
            json.dump(registry, f, indent=2)
    
    def add_algorithm(self,
                     source_file: str,
                     algorithm_info: Dict,
                     parameter_mapping: Dict,
                     category: str = 'custom',
                     metadata: Optional[Dict] = None) -> Dict[str, Any]:
        """
        Add a custom algorithm to the system.
        
        Parameters
        ----------
        source_file : str
            Path to source Python file
        algorithm_info : dict
            Parsed algorithm information
        parameter_mapping : dict
            Parameter mappings
        category : str
            Algorithm category
        metadata : dict, optional
            Additional metadata
            
        Returns
        -------
        result : dict
            Integration result
        """
        # Generate unique algorithm ID
        class_name = algorithm_info['main_class']['name']
        algo_id = f"custom_{class_name.lower()}"
        
        # Copy file to custom directory
        dest_file = self.custom_dir / f"{algo_id}.py"
        
        try:
            # Read source
            with open(source_file, 'r', encoding='utf-8') as f:
                source_code = f.read()
            
            # Add header comment
            header = f'''"""
Custom Algorithm: {class_name}
Uploaded: {metadata.get('upload_date', 'Unknown')}
Author: {metadata.get('author', 'Unknown')}
Category: {category}

This is a user-contributed algorithm integrated into MHA Flow.
"""

'''
            # Write to destination
            with open(dest_file, 'w', encoding='utf-8') as f:
                f.write(header + source_code)
            
            # Update registry
            registry = self._load_registry()
            registry['algorithms'][algo_id] = {
                'class_name': class_name,
                'file': str(dest_file.name),
                'category': category,
                'parameter_mapping': parameter_mapping,
                'metadata': metadata or {},
                'enabled': True
            }
            self._save_registry(registry)
            
            return {
                'success': True,
                'algorithm_id': algo_id,
                'file_path': str(dest_file),
                'message': f'Successfully integrated {class_name}'
            }
            
        except Exception as e:
            return {
                'success': False,
                'error': str(e),
                'message': f'Failed to integrate algorithm: {e}'
            }
    
    def load_custom_algorithm(self, algo_id: str):
        """
        Dynamically load a custom algorithm.
        
        Parameters
        ----------
        algo_id : str
            Algorithm ID
            
        Returns
        -------
        algorithm_class : class
            Loaded algorithm class
        """
        registry = self._load_registry()
        
        if algo_id not in registry['algorithms']:
            raise ValueError(f"Algorithm {algo_id} not found in registry")
        
        algo_info = registry['algorithms'][algo_id]
        file_path = self.custom_dir / algo_info['file']
        
        # Load module dynamically
        spec = importlib.util.spec_from_file_location(algo_id, file_path)
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        
        # Get the class
        class_name = algo_info['class_name']
        algorithm_class = getattr(module, class_name)
        
        return algorithm_class
    
    def list_custom_algorithms(self) -> List[Dict]:
        """List all registered custom algorithms"""
        registry = self._load_registry()
        return [
            {
                'id': algo_id,
                **algo_info
            }
            for algo_id, algo_info in registry['algorithms'].items()
            if algo_info.get('enabled', True)
        ]
    
    def remove_algorithm(self, algo_id: str) -> bool:
        """Remove a custom algorithm"""
        registry = self._load_registry()
        
        if algo_id in registry['algorithms']:
            # Delete file
            algo_info = registry['algorithms'][algo_id]
            file_path = self.custom_dir / algo_info['file']
            if file_path.exists():
                file_path.unlink()
            
            # Remove from registry
            del registry['algorithms'][algo_id]
            self._save_registry(registry)
            return True
        
        return False
    
    def _load_registry(self) -> Dict:
        """Load registry from file"""
        with open(self.registry_file, 'r') as f:
            return json.load(f)
    
    def _save_registry(self, registry: Dict):
        """Save registry to file"""
        with open(self.registry_file, 'w') as f:
            json.dump(registry, f, indent=2)


class CustomAlgorithmManager:
    """Main manager class orchestrating all custom algorithm operations"""
    
    def __init__(self):
        self.parser = AlgorithmParser()
        self.mapper = ParameterMapper()
        self.validator = AlgorithmValidator()
        self.integrator = AlgorithmIntegrator()
    
    def validate_algorithm_code(self, code_content: str) -> Dict[str, Any]:
        """
        Validate algorithm code from a string (e.g., uploaded file content).
        
        Parameters
        ----------
        code_content : str
            Python source code of the algorithm
            
        Returns
        -------
        result : dict
            Validation result with 'valid', 'class_name', 'errors', etc.
        """
        import ast
        import re
        
        result = {
            'valid': False,
            'class_name': None,
            'errors': [],
            'inherits_base': False,
            'parameters': {}
        }
        
        # Check syntax
        try:
            tree = ast.parse(code_content)
        except SyntaxError as e:
            result['errors'].append(f"Syntax error: {e}")
            return result
        
        # Find classes
        classes = [node for node in ast.walk(tree) if isinstance(node, ast.ClassDef)]
        if not classes:
            result['errors'].append("No class definition found in the file")
            return result
        
        # Find the main algorithm class (prefer one that inherits from BaseOptimizer)
        main_class = None
        for cls in classes:
            base_names = []
            for base in cls.bases:
                if isinstance(base, ast.Name):
                    base_names.append(base.id)
                elif isinstance(base, ast.Attribute):
                    base_names.append(base.attr)
            
            if any('Optimizer' in b or 'Base' in b for b in base_names):
                main_class = cls
                result['inherits_base'] = True
                break
        
        if main_class is None:
            # Use the first class
            main_class = classes[0]
        
        result['class_name'] = main_class.name
        
        # Check for required methods
        method_names = [node.name for node in ast.walk(main_class) if isinstance(node, ast.FunctionDef)]
        
        has_init = '__init__' in method_names
        has_optimize = 'optimize' in method_names or '_optimize' in method_names
        
        if not has_init:
            result['errors'].append("Missing __init__ method")
        
        if not has_optimize:
            result['errors'].append("Missing optimize or _optimize method")
        
        # Extract parameters from __init__
        for node in ast.walk(main_class):
            if isinstance(node, ast.FunctionDef) and node.name == '__init__':
                for arg in node.args.args:
                    if arg.arg != 'self':
                        result['parameters'][arg.arg] = 'any'
                break
        
        if not result['errors']:
            result['valid'] = True
        
        return result

    def process_upload(self, file_path: str, metadata: Optional[Dict] = None) -> Dict[str, Any]:
        """
        Complete workflow to process an uploaded algorithm file.
        
        Parameters
        ----------
        file_path : str
            Path to uploaded Python file
        metadata : dict, optional
            Additional metadata (author, description, etc.)
            
        Returns
        -------
        result : dict
            Processing result with status and next steps
        """
        result = {
            'stage': 'parsing',
            'success': False
        }
        
        # Stage 1: Parse file
        try:
            algorithm_info = self.parser.parse_file(file_path)
            result['parse_info'] = algorithm_info
            
            if 'error' in algorithm_info:
                result['error'] = algorithm_info['error']
                return result
                
        except Exception as e:
            result['error'] = f"Parsing failed: {e}"
            return result
        
        result['stage'] = 'validation'
        
        # Stage 2: Validate
        try:
            validation = self.validator.validate(algorithm_info)
            result['validation'] = validation
            
            if not validation['valid']:
                result['error'] = validation.get('error', 'Validation failed')
                result['warnings'] = validation.get('warnings', [])
                return result
                
        except Exception as e:
            result['error'] = f"Validation failed: {e}"
            return result
        
        result['stage'] = 'parameter_mapping'
        
        # Stage 3: Parameter mapping
        try:
            main_class = algorithm_info['main_class']
            init_method = next(m for m in main_class['methods'] if m['name'] == '__init__')
            param_names = [p['name'] for p in init_method['params'] if p['name'] != 'self']
            
            mapping_prompt = self.mapper.create_mapping_prompt(param_names)
            result['parameter_mapping'] = mapping_prompt
            
            # If unmapped parameters exist, need user input
            if mapping_prompt['requires_user_input']:
                result['needs_user_input'] = True
                result['message'] = 'Please confirm parameter mappings'
                return result
                
        except Exception as e:
            result['error'] = f"Parameter mapping failed: {e}"
            return result
        
        result['stage'] = 'ready_for_integration'
        result['success'] = True
        result['message'] = 'Algorithm validated and ready for integration'
        
        return result
    
    def integrate_algorithm(self,
                          file_path: str,
                          parameter_mapping: Dict,
                          category: str = 'custom',
                          metadata: Optional[Dict] = None) -> Dict[str, Any]:
        """
        Integrate a validated algorithm into the system.
        
        DEPRECATED: Use submit_to_github() instead for production.
        This method directly adds algorithms without admin review.
        
        Parameters
        ----------
        file_path : str
            Path to algorithm file
        parameter_mapping : dict
            Confirmed parameter mappings
        category : str
            Algorithm category
        metadata : dict, optional
            Additional metadata
            
        Returns
        -------
        result : dict
            Integration result
        """
        algorithm_info = self.parser.parse_file(file_path)
        
        return self.integrator.add_algorithm(
            file_path,
            algorithm_info,
            parameter_mapping,
            category,
            metadata
        )
    
    def submit_to_github(self,
                        file_path: str,
                        algorithm_name: str,
                        author: str,
                        description: str,
                        category: str = 'custom',
                        metadata: Optional[Dict] = None) -> Dict[str, Any]:
        """
        Submit custom algorithm to GitHub for admin review.
        
        This is the recommended way to add custom algorithms in production.
        Algorithm will be reviewed before being added to the main library.
        
        Parameters
        ----------
        file_path : str
            Path to algorithm file
        algorithm_name : str
            Name of the algorithm
        author : str
            Username/name of the submitter
        description : str
            Description of what the algorithm does
        category : str
            Algorithm category (e.g., 'swarm', 'evolutionary', 'hybrid')
        metadata : dict, optional
            Additional metadata
            
        Returns
        -------
        result : dict
            Submission result with GitHub PR information
        """
        if not self.use_github:
            return {
                'success': False,
                'error': 'GitHub integration not available. Please configure GITHUB_TOKEN and GITHUB_REPO environment variables.',
                'fallback': 'Contact administrator to manually add your algorithm.'
            }
        
        # Validate algorithm first
        validation_result = self.validate_algorithm(file_path)
        
        if not validation_result['success']:
            return {
                'success': False,
                'error': 'Algorithm failed validation',
                'details': validation_result
            }
        
        # Read algorithm file content
        with open(file_path, 'r', encoding='utf-8') as f:
            algorithm_content = f.read()
        
        # Prepare metadata
        full_metadata = {
            'algorithm_name': algorithm_name,
            'description': description,
            'category': category,
            'author': author,
            'syntax_valid': validation_result.get('syntax_valid', False),
            'structure_valid': validation_result['validation']['valid'],
            'inherits_base': True,
            'parameters': validation_result.get('parameter_mapping', {}).get('mapped_parameters', {}),
            'validation_details': validation_result['validation']['details']
        }
        
        if metadata:
            full_metadata.update(metadata)
        
        # Submit to GitHub
        try:
            result = self.github_uploader.submit_algorithm(
                algorithm_name=algorithm_name,
                algorithm_file_path=file_path,
                algorithm_content=algorithm_content,
                metadata=full_metadata,
                author=author
            )
            
            return result
            
        except Exception as e:
            return {
                'success': False,
                'error': f'GitHub submission failed: {str(e)}'
            }
    
    def check_submission_status(self, pr_number: int) -> Dict[str, Any]:
        """
        Check the status of a submitted algorithm.
        
        Parameters
        ----------
        pr_number : int
            Pull request number from the submission
            
        Returns
        -------
        status : dict
            Current status of the submission
        """
        if not self.use_github:
            return {
                'success': False,
                'error': 'GitHub integration not available'
            }
        
        return self.github_uploader.get_submission_status(pr_number)
    
    def list_custom_algorithms(self) -> List[Dict]:
        """List all registered custom algorithms"""
        return self.integrator.list_custom_algorithms()
    
    def remove_algorithm(self, algo_id: str) -> bool:
        """Remove a custom algorithm"""
        return self.integrator.remove_algorithm(algo_id)
    
    def get_algorithm(self, algo_id: str) -> Optional[Dict]:
        """Get information about a specific custom algorithm"""
        registry = self.integrator._load_registry()
        return registry['algorithms'].get(algo_id)


# Example usage and testing
if __name__ == "__main__":
    print("=" * 70)
    print("Custom Algorithm Manager - Test Suite")
    print("=" * 70)
    
    manager = CustomAlgorithmManager()
    
    # Test with a sample algorithm file (if exists)
    print("\n[Test] Algorithm parser and validator ready")
    print(f"Custom algorithms directory: {manager.integrator.custom_dir}")
    print(f"Registry file: {manager.integrator.registry_file}")
    
    # List any existing custom algorithms
    custom_algos = manager.list_custom_algorithms()
    print(f"\nRegistered custom algorithms: {len(custom_algos)}")
    for algo in custom_algos:
        print(f"  - {algo['id']}: {algo['class_name']} [{algo['category']}]")
    
    print("\n" + "=" * 70)
