# scilink/agents/sim_agents/lammps_analysis_updater.py

import os
import re
import logging
import json
from pathlib import Path
from typing import Dict, Any, Optional, List

from ...auth import get_internal_proxy_key
from ...wrappers.openai_wrapper import OpenAIAsGenerativeModel
from ...wrappers.litellm_wrapper import LiteLLMGenerativeModel
from ._deprecation import normalize_params


class LAMMPSAnalysisUpdater:
    """
    Updater that analyzes failed LAMMPS analysis code and generates corrections.
    
    Handles failures in analysis scripts generated by LAMMPSAnalysisAgent.
    """
    
    def __init__(self, 
                 api_key: Optional[str] = None, 
                 model_name: str = "gemini-3-pro-preview",
                 base_url: Optional[str] = None,
                 # Legacy parameters
                 local_model: Optional[str] = None,
                 google_api_key: Optional[str] = None):
        """
        Initialize the LAMMPS analysis updater.
        
        Args:
            api_key: API key for the LLM provider
            model_name: Model name to use
            base_url: Optional base URL for internal proxy
            local_model: Deprecated, use base_url instead
            google_api_key: Deprecated, use api_key instead
        """
        # Configure logging
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(logging.INFO)
        if not self.logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            self.logger.addHandler(handler)
        
        # Normalize deprecated parameters
        api_key, base_url = normalize_params(
            api_key=api_key,
            google_api_key=google_api_key,
            base_url=base_url,
            local_model=local_model,
            source="LAMMPSAnalysisUpdater"
        )
        
        # Initialize model
        if base_url:
            if api_key is None:
                api_key = get_internal_proxy_key()
            
            if not api_key:
                raise ValueError("API key required for internal proxy")
            
            self.logger.info(f"Using internal proxy: {base_url}")
            self.model = OpenAIAsGenerativeModel(
                model=model_name,
                api_key=api_key,
                base_url=base_url
            )
        else:
            self.logger.info(f"Using LiteLLM: {model_name}")
            self.model = LiteLLMGenerativeModel(
                model=model_name,
                api_key=api_key
            )
    
    # ============================================================================
    # HELPER METHODS FOR LLM CALLS
    # ============================================================================
    
    def _generate_json(self, prompt: str) -> Dict[str, Any]:
        """Generate JSON response from LLM."""
        try:
            response = self.model.generate_content(
                prompt,
                generation_config={"response_mime_type": "application/json"}
            )
            return json.loads(response.text)
        except json.JSONDecodeError as e:
            self.logger.error(f"Failed to parse JSON: {e}")
            text = response.text
            match = re.search(r'\{.*\}', text, re.DOTALL)
            if match:
                try:
                    return json.loads(match.group(0))
                except:
                    pass
            raise ValueError(f"Could not parse JSON: {e}")
    
    def _generate_text(self, prompt: str) -> str:
        """Generate text response from LLM."""
        try:
            response = self.model.generate_content(prompt)
            return response.text
        except Exception as e:
            self.logger.error(f"Error generating text: {e}")
            raise
    
    # ============================================================================
    # FAILURE ANALYSIS
    # ============================================================================
    
    def analyze_failure(self, 
                       analysis_name: str,
                       generated_code: str,
                       error_result: Dict[str, Any],
                       data_files: Dict[str, Any],
                       available_packages: List[str]) -> Dict[str, Any]:
        """
        Analyze why an analysis failed and suggest corrections.
        
        Args:
            analysis_name: Name of the failed analysis
            generated_code: The Python code that failed
            error_result: Error information from execution
            data_files: Available data files (can be paths or lists of paths)
            available_packages: List of available packages
            
        Returns:
            Analysis of the failure with suggested fixes
        """
        self.logger.info(f"Analyzing failure for: {analysis_name}")
        
        error_message = error_result.get('message', 'Unknown error')
        stderr = error_result.get('stderr', '')
        
        # Check if data_files is empty or malformed
        if not data_files:
            self.logger.error("❌ data_files dictionary is EMPTY!")
            return {
                "error_type": "configuration_error",
                "root_cause": "The data_files dictionary is empty. No data was provided for analysis.",
                "correction_strategy": "Verify output_data contains required data types and is correctly passed to refinement.",
                "code_issues": [],
                "data_file_issues": {},
                "missing_functionality": [],
                "needs_data_preprocessing": False,
                "suggested_packages_to_use": [],
                "data_type_mismatch": {}
            }
        
        # Try to peek at data files to understand their structure
        data_previews = {}
        for data_type, file_path_or_list in data_files.items():
            try:
                # Handle both single files and lists
                if isinstance(file_path_or_list, list):
                    sample_file = file_path_or_list[0] if file_path_or_list else None
                    if sample_file and os.path.exists(sample_file):
                        with open(sample_file, 'r') as f:
                            lines = f.readlines()[:20]
                            data_previews[data_type] = {
                                "type": "series",
                                "num_files": len(file_path_or_list),
                                "sample_preview": ''.join(lines),
                                "files": [os.path.basename(f) for f in file_path_or_list[:3]]
                            }
                    else:
                        data_previews[data_type] = {
                            "type": "series",
                            "num_files": len(file_path_or_list),
                            "error": "File not found or empty list"
                        }
                else:
                    # Single file
                    if os.path.exists(file_path_or_list):
                        with open(file_path_or_list, 'r') as f:
                            lines = f.readlines()[:20]
                            data_previews[data_type] = {
                                "type": "single",
                                "preview": ''.join(lines),
                                "file": os.path.basename(file_path_or_list)
                            }
                    else:
                        data_previews[data_type] = {
                            "type": "single",
                            "error": f"File not found: {file_path_or_list}"
                        }
            except Exception as e:
                self.logger.warning(f"Could not preview {data_type}: {e}")
                data_previews[data_type] = {"error": str(e)}
        
        analysis_prompt = f"""
You are an expert in LAMMPS simulation analysis and Python scientific computing.
An analysis script has failed. Analyze the error and suggest specific corrections.

ANALYSIS NAME: {analysis_name}

ERROR MESSAGE:
{error_message}

STDERR OUTPUT:
{stderr}

GENERATED CODE:
{generated_code}

AVAILABLE DATA FILES:
{json.dumps(data_files, indent=2)}

DATA FILE PREVIEWS (first 20 lines):
{json.dumps(data_previews, indent=2)}

AVAILABLE PACKAGES:
{', '.join(available_packages)}

🚨 COMMON ISSUES TO CHECK:
1. Empty data_files dictionary (configuration error)
2. Incorrect data type keys (e.g., expecting 'positions_trj' but receiving 'trajectory')
3. Not handling list vs string for file paths
4. File format issues (headers, delimiters, comments)
5. Missing file path validation
6. Timeout due to inefficient loops
7. Missing topology file for trajectory analysis
8. Incorrect MDAnalysis format specification (use LAMMPSDUMP not LAMMPSTRJ)

🚨 PATH REQUIREMENTS:
- ALL file saves must use: os.path.join(output_dir, "filename.ext")
- NEVER use relative paths like "./file.png" or "file.csv"

🚨 LAMMPS TRAJECTORY ANALYSIS REQUIREMENTS:
- LAMMPS trajectory files (.lammpstrj) require format='LAMMPSDUMP'
- Trajectory analysis requires BOTH topology (.data) AND trajectory (.lammpstrj)
- Example: u = mda.Universe('system.data', 'traj.lammpstrj', format='LAMMPSDUMP', topology_format='DATA')

Analyze the failure and respond with JSON:
{{
    "error_type": "configuration_error|data_format|package_missing|logic_error|file_not_found|type_error|timeout|topology_missing",
    "root_cause": "Detailed explanation of what went wrong",
    "data_file_issues": {{
        "has_header": true/false,
        "delimiter": "space|comma|tab",
        "column_count": 0,
        "data_starts_at_line": 0,
        "issues": ["issue1", "issue2"],
        "is_series": true/false,
        "needs_topology": true/false
    }},
    "code_issues": [
        {{
            "line_number": 0,
            "issue": "Description of the problem",
            "suggested_fix": "How to fix it"
        }}
    ],
    "missing_functionality": ["What's missing from the code"],
    "correction_strategy": "Overall approach to fix the code",
    "needs_data_preprocessing": true/false,
    "suggested_packages_to_use": ["package1", "package2"],
    "data_type_mismatch": {{
        "expected_keys": ["key1", "key2"],
        "actual_keys": ["key3", "key4"],
        "mapping_needed": true/false
    }}
}}

Return ONLY the JSON with no additional text.
"""
        
        try:
            analysis = self._generate_json(analysis_prompt)
            self.logger.info(f"Failure analysis completed: {analysis.get('error_type', 'unknown')}")
            return analysis
        except Exception as e:
            self.logger.error(f"Error during failure analysis: {e}")
            return {
                "error_type": "analysis_failed",
                "root_cause": f"Could not analyze failure: {str(e)}",
                "correction_strategy": "Manual review needed",
                "code_issues": [],
                "data_file_issues": {},
                "missing_functionality": [],
                "needs_data_preprocessing": False,
                "suggested_packages_to_use": [],
                "data_type_mismatch": {}
            }
    
    # ============================================================================
    # CODE CORRECTION
    # ============================================================================
    
    def generate_corrected_code(self,
                               analysis_name: str,
                               original_code: str,
                               failure_analysis: Dict[str, Any],
                               data_files: Dict[str, Any],
                               available_packages: List[str],
                               analysis_description: str = "") -> str:
        """
        Generate corrected analysis code based on failure analysis.
        
        Args:
            analysis_name: Name of the analysis
            original_code: Original (failed) code
            failure_analysis: Analysis from analyze_failure()
            data_files: Available data files
            available_packages: List of available packages
            analysis_description: Description of analysis purpose
            
        Returns:
            Corrected Python code
        """
        self.logger.info(f"Generating corrected code for: {analysis_name}")
        
        correction_prompt = f"""
Generate a corrected version of this LAMMPS analysis code.

ANALYSIS: {analysis_name}
DESCRIPTION: {analysis_description}

FAILURE ANALYSIS:
- Error Type: {failure_analysis.get('error_type', 'unknown')}
- Root Cause: {failure_analysis.get('root_cause', 'unknown')}
- Correction Strategy: {failure_analysis.get('correction_strategy', 'unknown')}

CODE ISSUES:
{json.dumps(failure_analysis.get('code_issues', []), indent=2)}

DATA FILE INFORMATION:
{json.dumps(failure_analysis.get('data_file_issues', {}), indent=2)}

DATA TYPE MISMATCH (if any):
{json.dumps(failure_analysis.get('data_type_mismatch', {}), indent=2)}

AVAILABLE DATA FILES:
{json.dumps(data_files, indent=2)}

ORIGINAL CODE (THAT FAILED):
{original_code}

AVAILABLE PACKAGES (YOU MUST USE ONLY THESE):
{', '.join(available_packages)}

═══════════════════════════════════════════════════════════════
🚨 CRITICAL REQUIREMENTS 🚨
═══════════════════════════════════════════════════════════════

1. DATA FILE HANDLING:
   - Check if data_files[key] is a LIST or STRING
   - For lists: data_files['type'] = [file1, file2, ...]
   - For single: data_files['type'] = "file_path"
   - Always validate: if not data_files: return {{'status': 'error'}}
   - Validate keys exist: if 'key' not in data_files: return error

2. FILE FORMAT HANDLING (LAMMPS-specific):
   - LAMMPS log files: may have comment lines starting with #
   - Thermodynamic data: usually space-separated
   - Skip header/comment lines before parsing
   - Check for empty data after filtering

3. TRAJECTORY ANALYSIS (CRITICAL):
   - LAMMPS format is 'LAMMPSDUMP' NOT 'LAMMPSTRJ'
   - Require topology (.data) file: u = mda.Universe(topology, traj, format='LAMMPSDUMP', topology_format='DATA')
   - Check both files exist before calling Universe

4. PATH REQUIREMENTS:
   - ALL outputs: output_path = os.path.join(output_dir, "file.ext")
   - NEVER use relative paths ("./file.png", "file.csv")

5. PERFORMANCE:
   - Avoid Python loops over large arrays
   - Use vectorized numpy operations
   - For trajectories: consider stride or max_frames limit

6. ERROR HANDLING:
   - Validate all file paths: os.path.exists(path)
   - Check data_files keys before access
   - Use try/except for file I/O
   - Return {{'status': 'error', 'message': ...}} on failure

7. OUTPUT FORMAT:
   - Always: print(json.dumps(results))
   - Include: 'status': 'success' or 'error'
   - Use absolute paths in results

═══════════════════════════════════════════════════════════════

REQUIRED CODE STRUCTURE:

import numpy as np
import matplotlib.pyplot as plt
import json
import os

def main(data_files, output_dir):
    results = {{}}
    
    try:
        os.makedirs(output_dir, exist_ok=True)
        
        # VALIDATE INPUT
        if not data_files:
            return {{'status': 'error', 'message': 'No data files provided'}}
        
        # CHECK FOR REQUIRED DATA
        required_key = 'appropriate_key'  # Based on analysis
        
        if required_key not in data_files:
            return {{
                'status': 'error',
                'message': f'Missing {{required_key}}. Available: {{list(data_files.keys())}}'
            }}
        
        file_or_list = data_files[required_key]
        
        # Handle list vs single file
        if isinstance(file_or_list, list):
            # Validate all files exist
            for f in file_or_list:
                if not os.path.exists(f):
                    return {{'status': 'error', 'message': f'File not found: {{f}}'}}
            # Process series...
        else:
            # Validate single file
            if not os.path.exists(file_or_list):
                return {{'status': 'error', 'message': f'File not found: {{file_or_list}}'}}
            # Process single file...
        
        # PERFORM ANALYSIS
        # ... analysis code ...
        
        # SAVE OUTPUTS
        plot_path = os.path.join(output_dir, 'plot.png')
        plt.savefig(plot_path, dpi=150, bbox_inches='tight')
        plt.close()
        results['plot'] = plot_path
        
        results['status'] = 'success'
        return results
        
    except Exception as e:
        import traceback
        return {{
            'status': 'error',
            'message': str(e),
            'traceback': traceback.format_exc()
        }}

if __name__ == "__main__":
    data_files = DATA_FILES_PLACEHOLDER
    output_dir = "OUTPUT_DIR_PLACEHOLDER"
    results = main(data_files, output_dir)
    print(json.dumps(results))

Generate COMPLETE corrected code addressing ALL issues.
Return ONLY Python code with NO markdown.
"""
        
        try:
            corrected_code = self._generate_text(correction_prompt)
            
            # Clean markdown formatting
            corrected_code = re.sub(r'^```python\s*', '', corrected_code, flags=re.MULTILINE)
            corrected_code = re.sub(r'^```\s*', '', corrected_code, flags=re.MULTILINE)
            corrected_code = re.sub(r'\s*```$', '', corrected_code, flags=re.MULTILINE)
            corrected_code = corrected_code.strip()
            
            self.logger.info(f"Corrected code generated ({len(corrected_code)} characters)")
            return corrected_code
            
        except Exception as e:
            self.logger.error(f"Error generating corrected code: {e}")
            # Return original code as fallback
            return original_code
