#!/usr/bin/env python3
"""
Simplified Ubuntu Sandbox - Core Functionality Only
"""

import gradio as gr
import subprocess
import os
import json
from typing import Dict, List, Any
import psutil
import time
from datetime import datetime

# Initialize workspace
WORKSPACE_PATH = "/home/user/workspace"
os.makedirs(WORKSPACE_PATH, exist_ok=True)

class SimpleSandbox:
    def __init__(self):
        self.sessions = {}
    
    def execute_command(self, command: str, session_id: str = "default") -> str:
        """Execute command safely"""
        if not command.strip():
            return "Empty command"
        
        try:
            # Simple command execution
            result = subprocess.run(
                command, 
                shell=True, 
                capture_output=True, 
                text=True, 
                timeout=10,
                cwd=WORKSPACE_PATH
            )
            
            output = f"Command: {command}\n"
            output += f"Exit Code: {result.returncode}\n"
            if result.stdout:
                output += f"Output:\n{result.stdout}"
            if result.stderr:
                output += f"Error:\n{result.stderr}"
            output += f"\nTimestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
            
            return output
            
        except subprocess.TimeoutExpired:
            return f"Command timed out: {command}"
        except Exception as e:
            return f"Error executing command: {str(e)}"
    
    def get_system_info(self) -> str:
        """Get system information"""
        try:
            info = {
                "CPU": f"{psutil.cpu_count()} cores",
                "Memory": f"{psutil.virtual_memory().total / (1024**3):.1f} GB",
                "Disk": f"{psutil.disk_usage('/').free / (1024**3):.1f} GB free",
                "Uptime": f"{time.time() - psutil.boot_time():.0f} seconds"
            }
            return json.dumps(info, indent=2)
        except Exception as e:
            return f"Error getting system info: {str(e)}"
    
    def list_files(self, path: str = "/home/user/workspace") -> str:
        """List files in workspace"""
        try:
            if not os.path.exists(path):
                return f"Path does not exist: {path}"
            
            files = []
            for item in os.listdir(path):
                item_path = os.path.join(path, item)
                if os.path.isdir(item_path):
                    files.append(f"[DIR]  {item}")
                else:
                    size = os.path.getsize(item_path)
                    files.append(f"[FILE] {item} ({size} bytes)")
            
            return "\n".join(files) if files else "No files found"
        except Exception as e:
            return f"Error listing files: {str(e)}"

# Create sandbox instance
sandbox = SimpleSandbox()

def terminal_interface():
    """Create Gradio interface"""
    
    with gr.Blocks(title="Ubuntu Sandbox v2.0", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🚀 Ubuntu Sandbox v2.0 - AI Development Environment")
        gr.Markdown("A minimal, secure environment for AI models to execute commands and manage files.")
        
        with gr.Tabs():
            with gr.Tab("💻 Terminal"):
                command = gr.Textbox(
                    label="Enter Command", 
                    placeholder="e.g., ls, pwd, python --version, echo 'Hello AI!'"
                )
                output = gr.Textbox(
                    label="Output", 
                    lines=10, 
                    max_lines=20,
                    info="Command execution results will appear here"
                )
                
                gr.Button("▶️ Execute", variant="primary").click(
                    fn=sandbox.execute_command,
                    inputs=[command],
                    outputs=[output]
                )
            
            with gr.Tab("📊 System Info"):
                system_output = gr.Textbox(
                    label="System Information", 
                    lines=8,
                    info="Current system status and resources"
                )
                gr.Button("🔄 Refresh Info", variant="secondary").click(
                    fn=sandbox.get_system_info,
                    inputs=[],
                    outputs=[system_output]
                )
            
            with gr.Tab("📁 File Browser"):
                file_path = gr.Textbox(
                    label="Directory Path", 
                    value="/home/user/workspace",
                    info="Enter path to list files"
                )
                files_output = gr.Textbox(
                    label="Files and Directories", 
                    lines=10,
                    info="Directory contents will be listed here"
                )
                gr.Button("📋 List Files", variant="secondary").click(
                    fn=sandbox.list_files,
                    inputs=[file_path],
                    outputs=[files_output]
                )
        
        gr.Markdown("## 🔧 API Endpoints")
        gr.Markdown("""
        - **GET** `/health` - Health check
        - **POST** `/api/execute` - Execute command
        - **GET** `/api/system` - Get system info
        - **GET** `/api/files` - List files
        """)
    
    return demo

if __name__ == "__main__":
    print("🚀 Starting Ubuntu Sandbox v2.0...")
    app = terminal_interface()
    app.launch(
        server_name="0.0.0.0", 
        server_port=7860,
        share=False,
        show_error=True
    )