import spaces
import subprocess
import os
import gradio as gr
from smolagents import CodeAgent, InferenceClientModel
from smolagents.tools import Tool

class CloudTerminalTool(Tool):
    name = "terminal"
    description = "Executes bash shell commands directly in the Linux container environment."
    inputs = {"command": {"type": "string", "description": "The exact shell command to execute."}}
    output_type = "string"

    def forward(self, command: str) -> str:
        try:
            res = subprocess.run(
                command, 
                shell=True, 
                capture_output=True, 
                text=True, 
                timeout=120
            )
            out = res.stdout if res.stdout else ""
            err = f"\n[stderr]\n{res.stderr}" if res.stderr else ""
            return (out + err).strip() if (out or err) else "Command executed successfully (no output)."
        except subprocess.TimeoutExpired:
            return "Error: Command timed out after 120 seconds."
        except Exception as e:
            return f"Execution error: {str(e)}"

# Using a standard free inference model client configuration
model = InferenceClientModel(
    model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
    token=os.getenv("HF_TOKEN")
)

agent = CodeAgent(tools=[CloudTerminalTool()], model=model)

@spaces.GPU
def run_agent_task(prompt):
    if not prompt.strip():
        return "Please enter a task or command."
    try:
        result = agent.run(prompt)
        return str(result)
    except Exception as e:
        return f"Agent Runtime Error:\n{str(e)}"

with gr.Blocks(title="Cloud Agent Terminal") as demo:
    gr.Markdown("## Free Cloud AI Agent with Terminal Access")
    
    with gr.Column():
        task_input = gr.Textbox(
            label="Agent Prompt / Task", 
            placeholder="e.g. Run 'uname -a' or list files",
            lines=2
        )
        run_btn = gr.Button("Execute Task", variant="primary")
        output_box = gr.Textbox(label="Terminal Output", lines=14, interactive=False)
    
    run_btn.click(
        fn=run_agent_task, 
        inputs=task_input, 
        outputs=output_box
    )

demo.launch()