import uuid
import matplotlib.pyplot as plt
from pathlib import Path
from typing import Dict, Any, List, Optional
import pandas as pd
import json
import io
import contextlib
import traceback
from pydantic import BaseModel

class CodeResponse(BaseModel):
    """Container for code-related responses"""
    language: str = "python"
    code: str

class ChartSpecification(BaseModel):
    """Details about requested charts"""
    image_description: str
    code: Optional[str] = None

class AnalysisOperation(BaseModel):
    """Container for a single analysis operation with its code and result"""
    code: CodeResponse
    description: str

class CsvChatResult(BaseModel):
    """Structured response for CSV-related AI interactions"""
    response_type: str  # Literal["casual", "data_analysis", "visualization", "mixed"]
    casual_response: str
    analysis_operations: List[AnalysisOperation]
    charts: Optional[List[ChartSpecification]] = None

class PythonExecutor:
    """Handles execution of Python code and dummy image generation for CSV analysis"""
    
    def __init__(self, df: pd.DataFrame, charts_folder: str = "charts"):
        """
        Initialize the PythonExecutor with a DataFrame
        
        Args:
            df (pd.DataFrame): The DataFrame to operate on
            charts_folder (str): Folder to save charts in
        """
        self.df = df
        self.charts_folder = Path(charts_folder)
        self.charts_folder.mkdir(exist_ok=True)
        
    def execute_code(self, code: str) -> Dict[str, Any]:
        """
        Execute Python code and return the output and any generated plots
        
        Args:
            code (str): Python code to execute
            
        Returns:
            dict: Dictionary containing execution results and any generated plots
        """
        output = ""
        error = None
        plots = []
        
        # Capture stdout
        stdout = io.StringIO()
        
        # Monkey patch plt.show() to save figures
        original_show = plt.show
        
        def custom_show():
            """Custom show function that saves plots instead of displaying them"""
            for i, fig in enumerate(plt.get_fignums()):
                figure = plt.figure(fig)
                # Save plot to bytes buffer
                buf = io.BytesIO()
                figure.savefig(buf, format='png', bbox_inches='tight')
                buf.seek(0)
                plots.append(buf.read())
            plt.close('all')
        
        try:
            # Create execution context with common libraries and the DataFrame
            exec_globals = {
                'pd': pd,
                'plt': plt,
                'json': json,
                'df': self.df,  # Include the DataFrame in the execution context
                '__builtins__': __builtins__,
            }
            
            # Replace plt.show with custom implementation
            plt.show = custom_show
            
            # Execute code and capture output
            with contextlib.redirect_stdout(stdout):
                exec(code, exec_globals)
            
            output = stdout.getvalue()
            
        except Exception as e:
            error = {
                "message": str(e),
                "traceback": traceback.format_exc()
            }
        finally:
            # Restore original plt.show
            plt.show = original_show
            
        return {
            'output': output,
            'error': error,
            'plots': plots
        }
    
    def save_plot_dummy(self, plot_data: bytes, description: str) -> str:
        """
        Save plot to charts folder and return a dummy URL
        
        Args:
            plot_data (bytes): Image data in bytes
            description (str): Description of the plot
            
        Returns:
            str: Dummy URL for the chart
        """
        # Generate unique filename
        filename = f"chart_{uuid.uuid4().hex}.png"
        filepath = self.charts_folder / filename
        
        # Save the plot (even though we're using dummy URLs, we still save it)
        with open(filepath, 'wb') as f:
            f.write(plot_data)
            
        # Return a dummy URL
        return f"https://example.com/charts/{filename}"
    
    def process_response(self, response: CsvChatResult) -> str:
        """
        Process the CsvChatResult response and generate formatted output
        
        Args:
            response (CsvChatResult): Response from CSV analysis
            
        Returns:
            str: Formatted output with results and dummy image URLs
        """
        output_parts = []
        
        # Add casual response
        output_parts.append(response.casual_response)
        
        # Process analysis operations
        for operation in response.analysis_operations:
            # Execute the code
            result = self.execute_code(operation.code.code)
            
            # Add operation description
            output_parts.append(f"\n{operation.description}:")
            
            # Add output or error
            if result['error']:
                output_parts.append(f"Error: {result['error']['message']}")
            else:
                output_parts.append(result['output'].strip())
        
        # Process charts if they exist
        if response.charts:
            output_parts.append("\nVisualizations:")
            
            for chart in response.charts:
                if chart.code:
                    # Execute the chart code
                    result = self.execute_code(chart.code)
                    
                    if result['plots']:
                        # Save each generated plot and get dummy URL
                        for plot_data in result['plots']:
                            dummy_url = self.save_plot_dummy(plot_data, chart.image_description)
                            output_parts.append(f"\n{chart.image_description}")
                            output_parts.append(f"![{chart.image_description}]({dummy_url})")
                    elif result['error']:
                        output_parts.append(f"\nError generating {chart.image_description}: {result['error']['message']}")
        
        return "\n".join(output_parts)