"""
Chain of Code (CoC) plugin for OptILLM.

This plugin implements a chain-of-code approach that combines Chain-of-Thought (CoT) 
reasoning with code execution and LLM-based code simulation.

SAFETY NOTE: This plugin has been refactored to use Jupyter notebook kernel execution
instead of direct exec() calls. This provides process isolation and prevents potentially
dangerous code from crashing or affecting the main OptILLM process.

Key safety improvements:
- Code runs in isolated notebook kernels (separate processes)
- 30-second timeout prevents infinite loops
- Main process is protected from crashes, system exits, and memory issues
- Matplotlib/visualization code is safely removed to prevent display issues
- Comprehensive error handling and recovery
"""

import re
import logging
from typing import Tuple, Dict, Any, List
import ast
import traceback
import math
import importlib
import json
import nbformat
from nbconvert.preprocessors import ExecutePreprocessor
import os
import tempfile

logger = logging.getLogger(__name__)

# Plugin identifier
SLUG = "coc"

# Maximum attempts to fix code
MAX_FIX_ATTEMPTS = 3

# Initial code generation prompt
CHAIN_OF_CODE_PROMPT = '''
Write Python code to solve this problem. The code should:
1. Break down the problem into clear computational steps
2. Use standard Python features and math operations
3. Store the final result in a variable named 'answer'
4. Include error handling where appropriate
5. Be complete and executable

Format your response using:
```python
[Your complete Python program here]
```
'''

# Code fix prompt
CODE_FIX_PROMPT = '''
The following Python code failed to execute. Fix the code to make it work.
Original code:
```python
{code}
```

Error encountered:
{error}

Please provide a complete, fixed version of the code that:
1. Addresses the error message
2. Maintains the same logic and approach
3. Stores the final result in 'answer'
4. Is complete and executable

Return only the fixed code in a code block:
```python
[Your fixed code here]
```
'''

# Simulation prompt
SIMULATION_PROMPT = '''
The following Python code could not be executed directly. Analyze the code and determine what the answer would be.
Pay special attention to:
1. The core computational logic, ignoring any visualization or display code
2. The key mathematical operations that determine the final answer
3. Any logic that affects the 'answer' variable

Code to analyze:
```python
{code}
```

Runtime error encountered:
{error}

Return ONLY the final value that would be in the 'answer' variable. Return just the value, no explanations.
'''

def extract_code_blocks(text: str) -> List[str]:
    """Extract Python code blocks from text."""
    pattern = r'```python\s*(.*?)\s*```'
    matches = re.findall(pattern, text, re.DOTALL)
    blocks = [m.strip() for m in matches]
    logger.info(f"Extracted {len(blocks)} code blocks")
    for i, block in enumerate(blocks):
        logger.info(f"Code block {i+1}:\n{block}")
    return blocks

def sanitize_code(code: str) -> str:
    """Prepare code for safe execution by removing problematic visualization code."""
    # Remove or modify problematic visualization code
    lines = code.split('\n')
    safe_lines = []
    for line in lines:
        # Skip matplotlib-related imports and plotting commands that could cause issues
        if any(x in line.lower() for x in ['matplotlib', 'plt.', '.plot(', '.show(', 'figure', 'subplot']):
            # Replace with a comment to maintain code structure
            safe_lines.append(f"# {line}  # Removed for safety")
        else:
            # Keep the line if it's not visualization-related
            safe_lines.append(line)
    
    return '\n'.join(safe_lines)

def execute_code(code: str) -> Tuple[Any, str]:
    """Attempt to execute the code using Jupyter notebook kernel and return result or error."""
    logger.info("Attempting to execute code in notebook kernel")
    logger.info(f"Code:\n{code}")
    
    try:
        # Sanitize the code first
        sanitized_code = sanitize_code(code)
        
        # Create a notebook with the code
        notebook = nbformat.v4.new_notebook()
        
        # Add code that captures the answer variable
        enhanced_code = f"""
{sanitized_code}

# Capture the answer variable for output
if 'answer' in locals():
    print(f"ANSWER_RESULT: {{answer}}")
else:
    print("ANSWER_RESULT: No answer variable found")
"""
        
        notebook['cells'] = [nbformat.v4.new_code_cell(enhanced_code)]
        
        # Convert notebook to JSON string and then to bytes
        notebook_json = nbformat.writes(notebook)
        notebook_bytes = notebook_json.encode('utf-8')

        # Create temporary notebook file
        with tempfile.NamedTemporaryFile(mode='wb', suffix='.ipynb', delete=False) as tmp:
            tmp.write(notebook_bytes)
            tmp.flush()
            tmp_name = tmp.name

        try:
            # Read and execute the notebook
            with open(tmp_name, 'r', encoding='utf-8') as f:
                nb = nbformat.read(f, as_version=4)
            
            # Execute with timeout and isolation
            ep = ExecutePreprocessor(timeout=30, kernel_name='python3')
            ep.preprocess(nb, {'metadata': {'path': './'}})

            # Extract the output
            output = ""
            error_output = ""
            
            for cell in nb.cells:
                if cell.cell_type == 'code' and cell.outputs:
                    for output_item in cell.outputs:
                        if output_item.output_type == 'stream':
                            if output_item.name == 'stdout':
                                output += output_item.text
                            elif output_item.name == 'stderr':
                                error_output += output_item.text
                        elif output_item.output_type == 'execute_result':
                            output += str(output_item.data.get('text/plain', ''))
                        elif output_item.output_type == 'error':
                            error_output += f"{output_item.ename}: {output_item.evalue}"
            
            # Check for errors first
            if error_output:
                logger.error(f"Execution failed: {error_output}")
                return None, error_output
            
            # Parse the answer from output
            output = output.strip()
            
            # Look for our special ANSWER_RESULT marker
            if "ANSWER_RESULT:" in output:
                answer_line = [line for line in output.split('\n') if 'ANSWER_RESULT:' in line][-1]
                answer_str = answer_line.split('ANSWER_RESULT:', 1)[1].strip()
                
                if answer_str == "No answer variable found":
                    error = "Code executed but did not produce an answer variable"
                    logger.warning(error)
                    return None, error
                
                try:
                    # Try to evaluate the answer to convert it to proper type
                    answer = ast.literal_eval(answer_str)
                except (ValueError, SyntaxError):
                    # If literal_eval fails, keep as string
                    answer = answer_str
                
                logger.info(f"Execution successful. Answer: {answer}")
                return answer, None
            else:
                # Fallback: try to extract answer from any output
                if output:
                    logger.info(f"Execution completed with output: {output}")
                    return output, None
                else:
                    error = "Code executed but produced no output"
                    logger.warning(error)
                    return None, error
                    
        finally:
            # Clean up temporary file
            try:
                os.unlink(tmp_name)
            except:
                pass
            
    except Exception as e:
        error = f"Notebook execution failed: {str(e)}"
        logger.error(error)
        return None, error

def generate_fixed_code(original_code: str, error: str, client, model: str) -> Tuple[str, int]:
    """Ask LLM to fix the broken code."""
    logger.info("Requesting code fix from LLM")
    logger.info(f"Original error: {error}")
    
    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": CODE_FIX_PROMPT.format(
                code=original_code, error=error)},
            {"role": "user", "content": "Fix the code to make it work."}
        ],
        temperature=0.2
    )
    
    fixed_code = response.choices[0].message.content
    code_blocks = extract_code_blocks(fixed_code)
    
    if code_blocks:
        logger.info("Received fixed code from LLM")
        return code_blocks[0], response.usage.completion_tokens
    else:
        logger.warning("No code block found in LLM response")
        return None, response.usage.completion_tokens

def simulate_execution(code: str, error: str, client, model: str) -> Tuple[Any, int]:
    """Ask LLM to simulate code execution."""
    logger.info("Attempting code simulation with LLM")
    
    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": SIMULATION_PROMPT.format(
                code=code, error=error)},
            {"role": "user", "content": "Simulate this code and return the final answer value."}
        ],
        temperature=0.2
    )
    
    try:
        result = response.choices[0].message.content.strip()
        # Try to convert to appropriate type
        try:
            answer = ast.literal_eval(result)
        except:
            answer = result
        logger.info(f"Simulation successful. Result: {answer}")
        return answer, response.usage.completion_tokens
    except Exception as e:
        logger.error(f"Failed to parse simulation result: {str(e)}")
        return None, response.usage.completion_tokens

def run(system_prompt: str, initial_query: str, client, model: str) -> Tuple[str, int]:
    """Main Chain of Code execution function."""
    logger.info("Starting Chain of Code execution")
    logger.info(f"Query: {initial_query}")
    
    # Initial code generation
    messages = [
        {"role": "system", "content": system_prompt + "\n" + CHAIN_OF_CODE_PROMPT},
        {"role": "user", "content": initial_query}
    ]
    
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0.7
    )
    total_tokens = response.usage.completion_tokens
    
    # Extract initial code
    code_blocks = extract_code_blocks(response.choices[0].message.content)
    if not code_blocks:
        logger.warning("No code blocks found in response")
        return response.choices[0].message.content, total_tokens

    current_code = code_blocks[0]
    fix_attempts = 0
    last_error = None
    
    # Strategy 1: Direct execution and fix attempts
    while fix_attempts < MAX_FIX_ATTEMPTS:
        fix_attempts += 1
        logger.info(f"Execution attempt {fix_attempts}/{MAX_FIX_ATTEMPTS}")
        
        # Try to execute current code
        answer, error = execute_code(current_code)
        
        # If successful, return the answer
        if error is None:
            logger.info(f"Successful execution on attempt {fix_attempts}")
            return str(answer), total_tokens
            
        last_error = error
        
        # If we hit max attempts, break to try simulation
        if fix_attempts >= MAX_FIX_ATTEMPTS:
            logger.warning(f"Failed after {fix_attempts} fix attempts")
            break
            
        # Otherwise, try to get fixed code from LLM
        logger.info(f"Requesting code fix, attempt {fix_attempts}")
        fixed_code, fix_tokens = generate_fixed_code(current_code, error, client, model)
        total_tokens += fix_tokens
        
        if fixed_code:
            current_code = fixed_code
        else:
            logger.error("Failed to get fixed code from LLM")
            break
    
    # Strategy 2: If all execution attempts failed, try simulation
    logger.info("All execution attempts failed, trying simulation")
    simulated_answer, sim_tokens = simulate_execution(current_code, last_error, client, model)
    total_tokens += sim_tokens
    
    if simulated_answer is not None:
        logger.info("Successfully got answer from simulation")
        return str(simulated_answer), total_tokens
    
    # If we get here, everything failed
    logger.warning("All strategies failed")
    return f"Error: Could not solve problem after all attempts. Last error: {last_error}", total_tokens