import gradio as gr
import os
from smolagents import InferenceClientModel, CodeAgent, MCPClient

def create_agent(hf_token):
    if not hf_token:
        return "please provide a valid hugging face token to continue"
    
    try:
        # initialize the mcp client
        mcp_client = MCPClient(
            {"url": "https://enricollen-mcp-sentiment.hf.space/gradio_api/mcp/sse"}
        )
        tools = mcp_client.get_tools()
        
        # create model with the provided token
        model = InferenceClientModel(token=hf_token)
        
        # create agent with tools
        agent = CodeAgent(
            tools=[*tools], 
            model=model, 
            additional_authorized_imports=["json", "ast", "urllib", "base64"]
        )
        
        return agent
    except Exception as e:
        return f"error creating agent: {str(e)}"

def agent_response(message, history, agent_state):
    if not agent_state or isinstance(agent_state, str):
        return history + [["please provide a valid hugging face token in the settings tab first", None]], agent_state
    
    try:
        response = agent_state.run(message)
        return history + [[message, str(response)]], agent_state
    except Exception as e:
        return history + [[message, f"error: {str(e)}"]], agent_state

# create the interface
with gr.Blocks() as demo:
    agent_state = gr.State(None)
    
    gr.Markdown("""
    # agent with mcp tools
    
    this is a simple agent that uses mcp tools to answer questions.
    
    1. first enter your hugging face api token below and click "setup agent"
    2. after setup is successful, go to the chat tab to start interacting with the agent
    
    example: "analyze the sentiment of the following text 'this is awesome'"
    """)
    
    with gr.Tabs() as tabs:
        with gr.Tab("Settings", id=0):
            with gr.Row():
                hf_token = gr.Textbox(
                    label="hugging face api token",
                    type="password",
                    placeholder="enter your hf api token here",
                    info="required to use the agent"
                )
                setup_button = gr.Button("setup agent")
            status = gr.Textbox(label="status", interactive=False)
            
            def setup_agent(token):
                result = create_agent(token)
                if isinstance(result, str):
                    return result, None
                return "agent successfully initialized", result
            
            setup_button.click(
                setup_agent,
                [hf_token],
                [status, agent_state]
            )
        
        with gr.Tab("Chat", id=1):
            chatbot = gr.Chatbot()
            with gr.Row():
                msg = gr.Textbox(
                    label="message",
                    placeholder="ask something...",
                    lines=2
                )
                submit_btn = gr.Button("enter")
            clear = gr.Button("clear")
            
            def submit_fn(message, history, state):
                if message.strip() == "":
                    return history, state
                return agent_response(message, history, state)
            
            msg.submit(
                submit_fn, 
                [msg, chatbot, agent_state], 
                [chatbot, agent_state],
                queue=True
            ).then(
                lambda: "", None, msg  
            )
            
            submit_btn.click(
                submit_fn,
                [msg, chatbot, agent_state], 
                [chatbot, agent_state],
                queue=True
            ).then(
                lambda: "", None, msg  
            )
            
            clear.click(lambda: [], None, chatbot, queue=False)

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