# space.py
import gradio as gr
import asyncio
import base64
from agent_core import run_agent, AIONAgent
import json
import os

# Ensure /data volume exists for memory persistence
os.makedirs("/data", exist_ok=True)

async def agent_wrapper(objective: str):
    try:
        result = await run_agent(objective)
        # Convert base64 screenshot to image for display
        img_data = base64.b64decode(result["final_screenshot"])
        # Save to temp
        img_path = "/data/latest_screenshot.png"
        with open(img_path, "wb") as f:
            f.write(img_data)
        # Format memory as readable text
        memory_text = json.dumps(result["memory"], indent=2)[:5000]  # truncate for UI
        return img_path, memory_text, result["status"]
    except Exception as e:
        return None, f"Error: {str(e)}", "failed"

def sync_wrapper(objective):
    return asyncio.run(agent_wrapper(objective))

# Gradio UI
with gr.Blocks(title="AION - Anonymous Web Agent") as demo:
    gr.Markdown("""
    # 🕵️ AION – Fully Anonymous Multi-Modal AI Agent
    Powered by **Playwright** + **GROQ Vision** + **OpenRouter**  
    All traffic routed through Tor + rotating residential proxies.  
    Screenshots captured and stored ephemerally.
    """)
    with gr.Row():
        with gr.Column(scale=2):
            objective_input = gr.Textbox(
                label="Objective",
                placeholder="e.g., Go to Wikipedia, search for 'AI safety', extract first paragraph, and screenshot.",
                lines=3
            )
            run_btn = gr.Button("🚀 Execute Agent", variant="primary")
        with gr.Column(scale=1):
            status_output = gr.Textbox(label="Status", interactive=False)
    
    with gr.Row():
        screenshot_output = gr.Image(label="Final Screenshot", type="filepath")
        memory_output = gr.Textbox(label="Memory & Reflections", lines=15, interactive=False)
    
    run_btn.click(
        fn=sync_wrapper,
        inputs=[objective_input],
        outputs=[screenshot_output, memory_output, status_output]
    )
    
    gr.Markdown("""
    ### 🔧 Hardcoded Credentials (Active)
    - GROQ: 2 keys (rotated round-robin)
    - OpenRouter: 7 keys (rotated round-robin)
    - Tor SOCKS5: `127.0.0.1:9050` with control on `9051`
    """)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, debug=False)