import gradio as gr
from agent import Agent
from persistence import AgentPersistence

# Instantiate persistence for shared use
persistence = AgentPersistence()

def initialize_agent(name, persona, instruction, strategy):
    """Initialize a new agent."""
    agent = Agent(name, persistence)
    agent.persona = persona
    agent.instruction = instruction
    agent.strategy = strategy
    agent.save_state()
    return f"Agent '{name}' initialized successfully."

def execute_task(agent_name, task):
    """Execute a task with the specified agent."""
    agent = Agent(agent_name, persistence)
    if not agent.load_state():
        return f"Error: Agent '{agent_name}' not found."
    result = agent.execute(task)
    return result

def pause_agent(agent_name):
    """Pause the agent by saving its state."""
    agent = Agent(agent_name, persistence)
    if agent.save_state():
        return f"Agent '{agent_name}' paused successfully."
    else:
        return f"Error: Failed to pause agent '{agent_name}'."

def resume_agent(agent_name):
    """Resume the agent's state."""
    agent = Agent(agent_name, persistence)
    if agent.load_state():
        return f"Agent '{agent_name}' resumed successfully. Current task: {agent.task}"
    else:
        return f"Error: Failed to resume agent '{agent_name}'."

def list_agents():
    """List all saved agents."""
    saved_agents = persistence.list_saved_agents()
    if not saved_agents:
        return "No agents found."
    return "\n".join([f"{name} (Last updated: {timestamp})" for name, timestamp in saved_agents.items()])

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Intelligent Agent Management")

    with gr.Tab("Initialize Agent"):
        agent_name = gr.Textbox(label="Agent Name")
        persona = gr.Textbox(label="Persona")
        instruction = gr.Textbox(label="Instruction")
        strategy = gr.Textbox(label="Strategy (e.g., ChainOfThoughtStrategy)")
        init_button = gr.Button("Initialize Agent")
        init_output = gr.Textbox(label="Output")
        init_button.click(
            initialize_agent,
            inputs=[agent_name, persona, instruction, strategy],
            outputs=init_output,
        )

    with gr.Tab("Execute Task"):
        exec_agent_name = gr.Textbox(label="Agent Name")
        task = gr.Textbox(label="Task")
        exec_button = gr.Button("Execute Task")
        exec_output = gr.Textbox(label="Output")
        exec_button.click(
            execute_task,
            inputs=[exec_agent_name, task],
            outputs=exec_output,
        )

    with gr.Tab("Pause Agent"):
        pause_agent_name = gr.Textbox(label="Agent Name")
        pause_button = gr.Button("Pause Agent")
        pause_output = gr.Textbox(label="Output")
        pause_button.click(
            pause_agent,
            inputs=[pause_agent_name],
            outputs=pause_output,
        )

    with gr.Tab("Resume Agent"):
        resume_agent_name = gr.Textbox(label="Agent Name")
        resume_button = gr.Button("Resume Agent")
        resume_output = gr.Textbox(label="Output")
        resume_button.click(
            resume_agent,
            inputs=[resume_agent_name],
            outputs=resume_output,
        )

    with gr.Tab("List Agents"):
        list_button = gr.Button("List All Agents")
        list_output = gr.Textbox(label="Agents")
        list_button.click(
            list_agents,
            inputs=[],
            outputs=list_output,
        )

demo.launch()