import datetime
import io
import time
from contextlib import redirect_stdout, redirect_stderr
import uuid
from matplotlib import pyplot as plt
import pandas as pd
import numpy as np
import traceback
import seaborn as sns


plt.style.use('seaborn-v0_8-whitegrid')

class PythonDataAnalysisExecutor:
    """
    Simplified Python code execution environment for data analysis
    """
    
    def __init__(self, timeout_seconds=30):
        self.timeout = timeout_seconds
        self.safe_globals = {
            '__builtins__': {
                'print': print, 'range': range, 'len': len,
                'str': str, 'int': int, 'float': float, 'bool': bool,
                'list': list, 'dict': dict, 'tuple': tuple, 'set': set,
                'min': min, 'max': max, 'sum': sum, 'abs': abs,
                'round': round, 'zip': zip, 'enumerate': enumerate,
                '__import__': __import__
            },
            'pd': pd, 'np': np,
            'matplotlib.pyplot': plt,
            'seaborn': sns,
            'uuid': uuid.uuid4,
            'datetime': datetime, 'time': time,
            'DataFrame': pd.DataFrame, 'Series': pd.Series
        }
        self.last_result = None
    
    def _validate_code(self, code):
        """Basic security checks"""
        forbidden = ['sys.', 'open(', 'exec(', 'eval(']
        if any(f in code for f in forbidden):
            raise ValueError("Potentially unsafe code detected")

    def execute(self, code: str) -> dict:
        """
        Execute Python code safely with timeout
        Returns dict with: success, output, error, execution_time, variables, result
        """
        result = {
            'success': False,
            'output': '',
            'error': '',
            'execution_time': 0,
            'variables': {},
            'result': None  # This will capture the last expression result
        }
        
        start_time = time.time()
        output = io.StringIO()
        
        try:
            self._validate_code(code)
            
            with redirect_stdout(output), redirect_stderr(output):
                # Split into lines and handle last expression
                lines = [line for line in code.split('\n') if line.strip()]
                if lines:
                    # Execute all but last line normally
                    if len(lines) > 1:
                        exec('\n'.join(lines[:-1]), self.safe_globals)
                    
                    # Handle last line specially to capture its value
                    last_line = lines[-1].strip()
                    if last_line:
                        # If it's an expression (not assignment or control structure)
                        if not (last_line.startswith((' ', '\t')) or 
                               last_line.split()[0] in ('if', 'for', 'while', 'def', 'class') or
                               '=' in last_line):
                            self.last_result = eval(last_line, self.safe_globals)
                            result['result'] = self.last_result
                            output.write(str(self.last_result) + '\n')
                        else:
                            exec(last_line, self.safe_globals)
            
            result['output'] = output.getvalue()
            result['variables'] = {
                k: v for k, v in self.safe_globals.items()
                if not k.startswith('__') and k in code
            }
            result['success'] = True
            
        except Exception as e:
            result['error'] = f"{str(e)}\n{traceback.format_exc()}"
        
        result['execution_time'] = time.time() - start_time
        return result


def run_analysis(code: str, timeout=20) -> dict:
    """Simplified interface for code execution"""
    executor = PythonDataAnalysisExecutor(timeout_seconds=timeout)
    return executor.execute(code)


# Example usage
# if __name__ == "__main__":
#     analysis_code = """
# import datetime
# print(datetime.datetime.now())
# """
    
#     result = run_analysis(analysis_code)
    
#     # Improved output formatting
#     if result['success']:
#         print("Execution successful")
#         print("Execution time:", result['execution_time'], "seconds")
#         print("Output:", result['output'].strip())
#         print("Result:", result['result'])
#         print("Variables:", list(result['variables'].keys()))
#     else:
#         print("Execution failed")
#         print("Error:", result['error'])