import gradio as gr
import json
import os
import PyPDF2

from mcp import StdioServerParameters
from smolagents import InferenceClientModel, CodeAgent, ToolCollection, MCPClient

def get_mcp_tools(mcp_servers: list[dict]) -> list[dict]:
    """
    Get MCP tools from MCP servers.

    Args:
        mcp_servers(list[dict]): List of MCP servers.

    Returns:
        list[dict]: List of MCP tools.
    """
    tools = []
    for mcp_server in mcp_servers:
        mcp_client = MCPClient(
            mcp_server
        )
        tools.extend(mcp_client.get_tools())
    return tools

MCP_SERVERS = [
    {"url": "https://abidlabs-mcp-tools2.hf.space/gradio_api/mcp/sse"},
    {"url": "https://proitm-mcpcoursesentimentanalysis.hf.space/gradio_api/mcp/sse"}
]
tools = get_mcp_tools(MCP_SERVERS)

model = InferenceClientModel(token=os.getenv("HF_TOKEN"))
# @TODO: Change promot to ask for clarifications if necessary and set final answer.
agent = CodeAgent(
    tools=[*tools],
    model=model,
    max_steps=4,
    planning_interval=3,
    stream_outputs=True,
)

memory = []

def format_message(memory, message):
    return "History: \n" + "\n".join([str(x) for x in memory]) + "\nTask: " + message

def extract_text_from_txt(file_obj):
    file_obj.seek(0)
    content = file_obj.read()
    if isinstance(content, bytes):
        try:
            content = content.decode("utf-8")
        except Exception:
            content = content.decode(errors="ignore")
    return content

def extract_text_from_pdf(file_obj):
    file_obj.seek(0)
    try:
        reader = PyPDF2.PdfReader(file_obj)
        pdf_text = ""
        for page in reader.pages:
            pdf_text += page.extract_text() or ""
        return pdf_text
    except Exception as e:
        return f"[Error reading PDF: {e}]"

def extract_text_from_file(file_obj):
    filename = getattr(file_obj, "name")
    if filename.lower().endswith(".txt"):
        return {"name": filename, "text": extract_text_from_txt(file_obj)}
    elif filename.lower().endswith(".pdf"):
        return {"name": filename, "text": extract_text_from_pdf(file_obj)}
    raise ValueError(f"Unsupported file type: {filename}")

def prepare_additional_args(message):
    additional_args = {}
    if "files" in message and message["files"]:
        file_texts = [extract_text_from_file(open(file_obj, "rb")) for file_obj in message["files"]]
        additional_args["uploaded_file_texts"] = file_texts
    return additional_args

def chat_function(message, history):
    # Add user message to memory
    memory.append({"message": message["text"], "role": "user", "type": "text"})

    # Prepare additional_args for the agent
    additional_args = prepare_additional_args(message)

    # Pass both the message and additional_args to the agent
    reply = agent.run(
        message["text"],
        additional_args=additional_args
    )

    memory.append({"message": reply, "role": "assistant", "type": "text"})
    return reply

chat = gr.ChatInterface(
    fn=chat_function,
    type="messages",
    examples=["Prime factorization of 68"],
    title="Agent with MCP Tools",
    description="This is a simple agent that uses MCP tools to answer questions.",
    multimodal=True,
    textbox=gr.MultimodalTextbox(file_types=[".pdf", ".txt"]),
    save_history=True,
)

chat.launch()