import gradio as gr
import subprocess
import os
import shlex

def run_shell_command(command, working_directory=""):
    """
    Execute a shell command and return the output.

    Args:
        command (str): The shell command to execute
        working_directory (str): Optional working directory to run the command in

    Returns:
        str: The command output or error message
    """
    if not command.strip():
        return "Error: Please enter a command"
    
    try:
        # Change to specified directory if provided
        original_cwd = os.getcwd()
        if working_directory and os.path.isdir(working_directory):
            os.chdir(working_directory)
        elif working_directory:
            return f"Error: Directory '{working_directory}' does not exist"
        
        # Execute the command with timeout for safety
        result = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=30  # 30 second timeout
        )
        
        # Restore original directory
        os.chdir(original_cwd)
        
        # Format output
        output = ""
        if result.stdout:
            output += f"STDOUT:\n{result.stdout}\n"
        if result.stderr:
            output += f"STDERR:\n{result.stderr}\n"
        
        output += f"Exit Code: {result.returncode}"
        
        return output if output else "Command executed successfully (no output)"
        
    except subprocess.TimeoutExpired:
        os.chdir(original_cwd)
        return "Error: Command timed out after 30 seconds"
    except Exception as e:
        os.chdir(original_cwd)
        return f"Error executing command: {str(e)}"

def get_current_directory():
    """Get the current working directory"""
    return os.getcwd()

# Create the main interface
with gr.Blocks(title="Shell Command Runner") as demo:
    gr.Markdown("# Shell Command Runner")
    gr.Markdown("Execute shell commands safely with timeout protection.")
    
    with gr.Row():
        with gr.Column():
            command_input = gr.Textbox(
                label="Command",
                placeholder="ls -la",
                lines=2
            )
            directory_input = gr.Textbox(
                label="Working Directory (optional)",
                placeholder="/path/to/directory",
                value=""
            )
            run_button = gr.Button("Run Command", variant="primary")
            
        with gr.Column():
            output_display = gr.Textbox(
                label="Output",
                lines=15,
                max_lines=20,
                interactive=False
            )
    
    with gr.Row():
        current_dir_button = gr.Button("Show Current Directory")
        current_dir_output = gr.Textbox(
            label="Current Directory",
            interactive=False
        )
    
    # Event handlers
    run_button.click(
        fn=run_shell_command,
        inputs=[command_input, directory_input],
        outputs=output_display
    )
    
    current_dir_button.click(
        fn=get_current_directory,
        outputs=current_dir_output
    )
    
    # Example commands
    gr.Markdown("""
    ### Example Commands:
    - `ls -la` - List files with details
    - `pwd` - Show current directory
    - `echo "Hello World"` - Print text
    - `date` - Show current date and time
    - `whoami` - Show current user
    - `df -h` - Show disk usage
    - `ps aux` - Show running processes
    
    **Note:** Commands timeout after 30 seconds for safety.
    """)

if __name__ == "__main__":
    demo.launch(mcp_server=True)