"""
AI-Parrot gVisor Sandbox Tool

Secure Python execution tool using gVisor (runsc) for complete kernel-level isolation.
This tool provides safe code execution for untrusted LLM-generated code.
"""
import json
import tempfile
import uuid
import shutil
import subprocess
import logging
from pathlib import Path
from typing import Any, Dict, Optional, List, Union
from dataclasses import dataclass, field
from contextlib import contextmanager
import base64
import pickle
import pandas as pd
from parrot.tools.abstract import AbstractTool, ToolResult


@dataclass
class SandboxConfig:
    """Configuration for gVisor sandbox"""
    runtime: str = "runsc"
    network: str = "none"  # Disable network by default
    max_memory: str = "2G"
    max_cpu: float = 2.0
    timeout: int = 30
    python_image: str = "python:3.11-slim"
    enable_gpu: bool = False
    mount_paths: List[str] = field(default_factory=list)
    pip_packages: List[str] = field(default_factory=lambda: [
        "pandas", "numpy", "matplotlib", "seaborn", "scikit-learn",
        "scipy", "plotly", "jupyterlab", "ipykernel"
    ])


@dataclass
class ExecutionResult:
    """Result from sandbox execution"""
    success: bool
    stdout: str
    stderr: str
    exit_code: int
    execution_time_ms: float
    files_created: Dict[str, bytes] = field(default_factory=dict)
    dataframes: Dict[str, pd.DataFrame] = field(default_factory=dict)
    plots: List[Dict[str, Any]] = field(default_factory=list)
    notebook_path: Optional[str] = None
    error_message: Optional[str] = None


class SandboxTool(AbstractTool):
    """
    Secure Python execution using gVisor sandbox.
    
    This tool provides kernel-level isolation for executing untrusted code
    generated by LLMs, preventing escape attempts and resource exhaustion.
    """
    
    name = "gvisor_python_sandbox"
    description = "Execute Python code in a secure gVisor sandbox with kernel isolation"
    
    def __init__(
        self,
        config: Optional[SandboxConfig] = None,
        workspace_dir: Optional[str] = None,
        dataframes: Optional[Dict[str, pd.DataFrame]] = None,
        enable_jupyter: bool = False,
        logger: Optional[logging.Logger] = None
    ):
        """
        Initialize gVisor sandbox tool.
        
        Args:
            config: Sandbox configuration
            workspace_dir: Directory for temporary files
            dataframes: Initial dataframes to make available
            enable_jupyter: Enable Jupyter notebook creation
            logger: Logger instance
        """
        super().__init__()
        self.config = config or SandboxConfig()
        self.workspace_dir = Path(workspace_dir or tempfile.mkdtemp(prefix="gvisor_"))
        self.workspace_dir.mkdir(parents=True, exist_ok=True)
        
        self.dataframes = dataframes or {}
        self.enable_jupyter = enable_jupyter
        self.logger = logger or logging.getLogger(__name__)
        
        # Verify gVisor installation
        self._verify_installation()
        
        # Prepare base container image
        self._prepare_base_image()
    
    def _verify_installation(self):
        """Verify gVisor and containerd are installed"""
        try:
            # Check runsc
            result = subprocess.run(
                ["runsc", "--version"],
                capture_output=True,
                text=True,
                check=False
            )
            if result.returncode != 0:
                self.logger.warning("runsc not found. Falling back to default docker runtime 'runc' (less secure).")
                self.config.runtime = "runc"
            else:
                self.logger.info(f"gVisor version: {result.stdout.strip()}")
        except FileNotFoundError:
            self.logger.warning("runsc not found. Falling back to default docker runtime 'runc' (less secure).")
            self.config.runtime = "runc"
            
        try:
            # Check containerd
            result = subprocess.run(
                ["ctr", "--version"],
                capture_output=True,
                text=True,
                check=False
            )
            if result.returncode != 0:
                self.logger.warning("containerd not found. Using Docker as fallback.")
                self.use_docker = True
            else:
                self.use_docker = False
        except FileNotFoundError:
            self.logger.warning("containerd not found. Using Docker as fallback.")
            self.use_docker = True
    
    def _prepare_base_image(self):
        """Prepare base container image with required packages"""
        dockerfile_content = f"""
FROM {self.config.python_image}

# Install system dependencies
RUN apt-get update && apt-get install -y \\
    gcc \\
    g++ \\
    make \\
    libssl-dev \\
    libffi-dev \\
    && rm -rf /var/lib/apt/lists/*

# Install Python packages
RUN pip install --no-cache-dir \\
    {' '.join(self.config.pip_packages)}

# Create working directory
WORKDIR /workspace

# Create sandbox user (non-root)
RUN useradd -m -s /bin/bash sandbox && \\
    chown -R sandbox:sandbox /workspace

USER sandbox

# Set Python to unbuffered mode
ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1
"""
        
        # Build image
        image_name = "ai-parrot-gvisor-sandbox"
        dockerfile_path = self.workspace_dir / "Dockerfile"
        dockerfile_path.write_text(dockerfile_content)
        
        try:
            subprocess.run(
                ["docker", "build", "-t", image_name, "-f", str(dockerfile_path), "."],
                cwd=self.workspace_dir,
                check=True,
                capture_output=True
            )
            self.image_name = image_name
            self.logger.info(f"Built sandbox image: {image_name}")
        except subprocess.CalledProcessError as e:
            self.logger.error(f"Failed to build image: {e.stderr.decode()}")
            raise
    
    @contextmanager
    def _create_sandbox_container(self, code: str, session_id: str):
        """Create and manage a gVisor sandbox container"""
        container_name = f"sandbox-{session_id}"
        exec_dir = self.workspace_dir / session_id
        exec_dir.mkdir(parents=True, exist_ok=True)
        
        try:
            # Prepare execution script
            script_path = exec_dir / "execute.py"
            script_path.write_text(self._prepare_execution_script(code))
            
            # Prepare data files
            self._prepare_data_files(exec_dir)
            
            # Create container with gVisor runtime
            run_cmd = [
                "docker", "run",
                "--runtime", self.config.runtime,  # Use gVisor runtime
                "--name", container_name,
                "-d",  # Detached
                "--rm",  # Auto-remove
                "--memory", self.config.max_memory,
                "--cpus", str(self.config.max_cpu),
                "--network", self.config.network,
                "--security-opt", "no-new-privileges",
                "--cap-drop", "ALL",  # Drop all capabilities
                "-v", f"{exec_dir}:/workspace:ro",  # Read-only mount
                "-v", f"{exec_dir}/output:/output:rw",  # Output directory
                self.image_name,
                "python", "/workspace/execute.py"
            ]
            
            # Start container
            subprocess.run(run_cmd, check=True, capture_output=True)
            
            yield container_name
            
        finally:
            # Cleanup container
            try:
                subprocess.run(
                    ["docker", "stop", container_name],
                    capture_output=True,
                    timeout=5,
                    check=False
                )
            except:
                subprocess.run(
                    ["docker", "kill", container_name],
                    capture_output=True,
                    check=False
                )
    
    def _prepare_execution_script(self, code: str) -> str:
        """Prepare the Python script for execution in sandbox"""
        return f'''
import sys
import os
import json
import pickle
import base64
import traceback
import pandas as pd
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Non-interactive backend
import matplotlib.pyplot as plt
from datetime import datetime

# Load serialized dataframes
dataframes_file = '/workspace/dataframes.pkl'
if os.path.exists(dataframes_file):
    with open(dataframes_file, 'rb') as f:
        dataframes = pickle.load(f)
        # Inject dataframes into global namespace
        for name, df in dataframes.items():
            globals()[name] = df
            print(f"Loaded dataframe: {{name}} with shape {{df.shape}}")
        if 'name' in locals():
            del name, df

# Execution metadata
execution_start = datetime.now()
created_files = {{}}
output_dataframes = {{}}
plots = []

# Redirect stdout to capture prints
import io
from contextlib import redirect_stdout, redirect_stderr

stdout_capture = io.StringIO()
stderr_capture = io.StringIO()

# Custom figure save wrapper
original_savefig = plt.savefig
def wrapped_savefig(fname, **kwargs):
    fname = f"/output/{{fname}}"
    original_savefig(fname, **kwargs)
    with open(fname, 'rb') as f:
        created_files[os.path.basename(fname)] = base64.b64encode(f.read()).decode()
    plots.append({{"type": "matplotlib", "filename": os.path.basename(fname)}})
    return fname
plt.savefig = wrapped_savefig

try:
    with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
        # Execute user code
        exec("""
{code}
""")
        
        # Capture any created dataframes
        for name, obj in list(globals().items()):
            if isinstance(obj, pd.DataFrame) and name not in ['dataframes', 'pd']:
                if name not in dataframes or not obj.equals(dataframes.get(name, pd.DataFrame())):
                    output_dataframes[name] = obj
    
    execution_time = (datetime.now() - execution_start).total_seconds() * 1000
    
    # Save results
    result = {{
        "success": True,
        "stdout": stdout_capture.getvalue(),
        "stderr": stderr_capture.getvalue(),
        "exit_code": 0,
        "execution_time_ms": execution_time,
        "files_created": created_files,
        "dataframes": {{name: df.to_dict('records') for name, df in output_dataframes.items()}},
        "plots": plots
    }}
    
except Exception as e:
    result = {{
        "success": False,
        "stdout": stdout_capture.getvalue(),
        "stderr": stderr_capture.getvalue() + "\\n" + traceback.format_exc(),
        "exit_code": 1,
        "execution_time_ms": (datetime.now() - execution_start).total_seconds() * 1000,
        "error_message": str(e)
    }}

# Write result
with open('/output/result.json', 'w') as f:
    json.dump(result, f)
'''
    
    def _prepare_data_files(self, exec_dir: Path):
        """Prepare data files for the sandbox"""
        # Create output directory
        output_dir = exec_dir / "output"
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # Serialize dataframes
        if self.dataframes:
            with open(exec_dir / "dataframes.pkl", 'wb') as f:
                pickle.dump(self.dataframes, f)
    
    async def _execute(
        self,
        code: str,
        description: str = "Execute Python code in gVisor sandbox",
        create_notebook: bool = False,
        **kwargs
    ) -> ToolResult:
        """
        Execute Python code in gVisor sandbox.
        
        Args:
            code: Python code to execute
            description: Description of execution
            create_notebook: Create Jupyter notebook of execution
            **kwargs: Additional parameters
            
        Returns:
            ToolResult with execution results
        """
        session_id = str(uuid.uuid4())[:8]
        self.logger.info(f"Starting sandbox execution: {session_id}")
        
        try:
            # Execute in sandbox
            result = await self._execute_in_sandbox(code, session_id)
            
            # Optionally create notebook
            if create_notebook and self.enable_jupyter:
                notebook_path = await self._create_notebook(code, result, session_id)
                result.notebook_path = notebook_path
            
            # Format response
            response = self._format_response(result)
            
            return ToolResult(
                result=response,
                status="success" if result.success else "error",
                error=result.error_message,
                metadata={
                    "session_id": session_id,
                    "execution_time_ms": result.execution_time_ms,
                    "files_created": list(result.files_created.keys()),
                    "dataframes_created": list(result.dataframes.keys()),
                    "notebook_path": result.notebook_path
                }
            )
            
        except Exception as e:
            self.logger.error(f"Sandbox execution failed: {e}")
            return ToolResult(
                result=None,
                status="error",
                error=str(e)
            )
        finally:
            # Cleanup session files
            if not kwargs.get("keep_files", False):
                self._cleanup_session(session_id)
    
    async def _execute_in_sandbox(self, code: str, session_id: str) -> ExecutionResult:
        """Execute code in gVisor sandbox container"""
        exec_dir = self.workspace_dir / session_id
        
        with self._create_sandbox_container(code, session_id) as container_name:
            # Wait for execution with timeout
            try:
                result = subprocess.run(
                    ["docker", "wait", container_name],
                    capture_output=True,
                    text=True,
                    timeout=self.config.timeout,
                    check=False
                )
                
                # Get logs
                logs = subprocess.run(
                    ["docker", "logs", container_name],
                    capture_output=True,
                    text=True,
                    check=False
                )
                
            except subprocess.TimeoutExpired:
                # Kill container on timeout
                subprocess.run(["docker", "kill", container_name], capture_output=True, check=False)
                return ExecutionResult(
                    success=False,
                    stdout="",
                    stderr="Execution timeout exceeded",
                    exit_code=-1,
                    execution_time_ms=self.config.timeout * 1000,
                    error_message=f"Code execution exceeded {self.config.timeout}s timeout"
                )
        
        # Read results
        result_file = exec_dir / "output" / "result.json"
        if result_file.exists():
            with open(result_file, 'r') as f:
                result_data = json.load(f)
            
            # Decode file contents
            files_created = {}
            for filename, content in result_data.get("files_created", {}).items():
                files_created[filename] = base64.b64decode(content)
            
            # Reconstruct dataframes
            dataframes = {}
            for name, records in result_data.get("dataframes", {}).items():
                dataframes[name] = pd.DataFrame(records)
            
            return ExecutionResult(
                success=result_data["success"],
                stdout=result_data["stdout"],
                stderr=result_data["stderr"],
                exit_code=result_data["exit_code"],
                execution_time_ms=result_data["execution_time_ms"],
                files_created=files_created,
                dataframes=dataframes,
                plots=result_data.get("plots", []),
                error_message=result_data.get("error_message")
            )
        else:
            return ExecutionResult(
                success=False,
                stdout=logs.stdout,
                stderr=logs.stderr,
                exit_code=result.returncode,
                execution_time_ms=0,
                error_message="Failed to read execution results"
            )
    
    async def _create_notebook(
        self,
        code: str,
        result: ExecutionResult,
        session_id: str
    ) -> str:
        """Create Jupyter notebook from execution"""
        import nbformat
        from nbformat.v4 import new_notebook, new_code_cell, new_markdown_cell
        
        # Create notebook
        nb = new_notebook()
        
        # Add header cell
        nb.cells.append(new_markdown_cell(
            f"# gVisor Sandbox Execution\n"
            f"Session: {session_id}\n"
            f"Execution Time: {result.execution_time_ms:.2f}ms"
        ))
        
        # Add code cell
        code_cell = new_code_cell(code)
        if result.stdout:
            code_cell.outputs = [{
                "output_type": "stream",
                "name": "stdout",
                "text": result.stdout
            }]
        nb.cells.append(code_cell)
        
        # Add results cell
        if result.dataframes:
            nb.cells.append(new_markdown_cell("## Created DataFrames"))
            for name, df in result.dataframes.items():
                nb.cells.append(new_code_cell(f"# DataFrame: {name}\n{df.to_string()}"))
        
        if result.plots:
            nb.cells.append(new_markdown_cell("## Generated Plots"))
            for plot in result.plots:
                nb.cells.append(new_markdown_cell(f"![{plot['filename']}]({plot['filename']})"))
        
        # Save notebook
        notebook_path = self.workspace_dir / f"notebook_{session_id}.ipynb"
        with open(notebook_path, 'w') as f:
            nbformat.write(nb, f)
        
        return str(notebook_path)
    
    def _format_response(self, result: ExecutionResult) -> str:
        """Format execution result for display"""
        parts = []
        
        if result.success:
            parts.append("✓ Code executed successfully")
            parts.append(f"Execution time: {result.execution_time_ms:.2f}ms")
        else:
            parts.append("✗ Code execution failed")
            parts.append(f"Error: {result.error_message}")
        
        if result.stdout:
            parts.append("\n--- Output ---")
            parts.append(result.stdout)
        
        if result.stderr and not result.success:
            parts.append("\n--- Errors ---")
            parts.append(result.stderr)
        
        if result.files_created:
            parts.append(f"\n--- Files Created ({len(result.files_created)}) ---")
            for filename in result.files_created.keys():
                parts.append(f"  • {filename}")
        
        if result.dataframes:
            parts.append(f"\n--- DataFrames Created ({len(result.dataframes)}) ---")
            for name, df in result.dataframes.items():
                parts.append(f"  • {name}: {df.shape[0]} rows × {df.shape[1]} columns")
        
        if result.notebook_path:
            parts.append(f"\n--- Notebook ---")
            parts.append(f"  📓 {result.notebook_path}")
        
        return "\n".join(parts)
    
    def _cleanup_session(self, session_id: str):
        """Clean up session files"""
        session_dir = self.workspace_dir / session_id
        if session_dir.exists():
            shutil.rmtree(session_dir)
    
    def add_dataframe(self, name: str, df: pd.DataFrame):
        """Add a dataframe to be available in the sandbox"""
        self.dataframes[name] = df
    
    def set_timeout(self, timeout: int):
        """Set execution timeout in seconds"""
        self.config.timeout = timeout
    
    def set_memory_limit(self, memory: str):
        """Set memory limit (e.g., '1G', '512M')"""
        self.config.max_memory = memory
    
    def cleanup(self):
        """Clean up all resources"""
        if self.workspace_dir.exists():
            shutil.rmtree(self.workspace_dir)


class SandboxPandasTool(SandboxTool):
    """
    Specialized version for Pandas operations with enhanced data handling.
    Drop-in replacement for PythonPandasTool with security.
    """
    
    name = "sandbox_pandas_tool"
    description = "Secure Pandas operations in gVisor sandbox"
    
    def __init__(self, dataframes: Optional[Dict[str, pd.DataFrame]] = None, **kwargs):
        """Initialize with Pandas-specific configuration"""
        config = SandboxConfig(
            pip_packages=[
                "pandas", "numpy", "matplotlib", "seaborn", 
                "plotly", "scikit-learn", "scipy", "statsmodels",
                "openpyxl", "xlrd", "pyarrow", "fastparquet"
            ]
        )
        super().__init__(config=config, dataframes=dataframes, **kwargs)
    
    async def analyze_dataframe(
        self,
        df_name: str,
        analysis_type: str = "summary"
    ) -> ToolResult:
        """
        Perform automated DataFrame analysis.
        
        Args:
            df_name: Name of dataframe to analyze
            analysis_type: Type of analysis (summary, correlation, distribution)
        """
        if df_name not in self.dataframes:
            return ToolResult(
                result=None,
                status="error", 
                error=f"DataFrame '{df_name}' not found"
            )
        
        analysis_code = {
            "summary": f"""
import pandas as pd
df = {df_name}
print("=== DataFrame Summary ===")
print(f"Shape: {{df.shape}}")
print(f"\\nData Types:\\n{{df.dtypes}}")
print(f"\\nSummary Statistics:\\n{{df.describe()}}")
print(f"\\nMissing Values:\\n{{df.isnull().sum()}}")
print(f"\\nMemory Usage:\\n{{df.memory_usage(deep=True)}}")
            """,
            
            "correlation": f"""
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

df = {df_name}
numeric_cols = df.select_dtypes(include=['number']).columns
if len(numeric_cols) > 1:
    corr_matrix = df[numeric_cols].corr()
    print("=== Correlation Matrix ===")
    print(corr_matrix)
    
    plt.figure(figsize=(10, 8))
    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', center=0)
    plt.title('Correlation Heatmap')
    plt.tight_layout()
    plt.savefig('correlation_heatmap.png')
else:
    print("Not enough numeric columns for correlation analysis")
            """,
            
            "distribution": f"""
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

df = {df_name}
numeric_cols = df.select_dtypes(include=['number']).columns

for i, col in enumerate(numeric_cols[:6]):  # Limit to first 6 columns
    plt.figure(figsize=(12, 4))
    
    plt.subplot(1, 3, 1)
    df[col].hist(bins=30, edgecolor='black')
    plt.title(f'{{col}} - Histogram')
    
    plt.subplot(1, 3, 2)
    df.boxplot(column=col)
    plt.title(f'{{col}} - Boxplot')
    
    plt.subplot(1, 3, 3)
    df[col].plot(kind='kde')
    plt.title(f'{{col}} - KDE')
    
    plt.tight_layout()
    plt.savefig(f'distribution_{{col}}.png')
    plt.close()
    
print(f"Generated distribution plots for {{len(numeric_cols)}} numeric columns")
            """
        }
        
        code = analysis_code.get(analysis_type, analysis_code["summary"])
        return await self.execute(
            code,
            description=f"Analyze DataFrame '{df_name}' - {analysis_type}"
        )


# Factory function for easy instantiation
def create_sandbox_tool(
    tool_type: str = "sandbox",
    **kwargs
) -> Union[SandboxTool, SandboxPandasTool]:
    """
    Factory function to create gVisor tools.
    
    Args:
        tool_type: Type of tool ('sandbox' or 'pandas')
        **kwargs: Configuration parameters
    
    Returns:
        Configured gVisor tool instance
    """
    if tool_type == "pandas":
        return SandboxPandasTool(**kwargs)
    else:
        return SandboxTool(**kwargs)