import asyncio
from langchain.agents import create_react_agent, AgentExecutor
from langchain.chat_models import init_chat_model
from langchain_core.prompts import PromptTemplate
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_core.tools import tool
from langchain.memory import MongoDBChatMessageHistory, ConversationBufferWindowMemory

from agents.agent_templates import AgentTemplate
from agents.output_parser.MCPReactParser import MCPReactParser


@tool
def get_daily_string(salt: str):
    """
    :param salt: a random string
    :return:get today's string
    """
    return "assss"


async def main():

    temp1 = AgentTemplate()
    temp1.set_model("gpt-4.1-nano", "openai", "")
    temp1.set_mcp_server("terminal_tools", {
        'transport': 'sse',
        'headers': {
            "X-Api-Key": "bd303c2f2a515b4ce70c1f07ee85530c2"
        },
        'url': 'http://localhost:5050/sse'
    })


    mcp_client = MultiServerMCPClient({
        'terminal_tool':{
            'transport': 'sse',
            'headers': {
                "X-Api-Key": "bd303c2f2a515b4ce70c1f07ee85530c2"
            },
            'url': 'http://localhost:5050/sse'
        }
    })

    tools = await mcp_client.get_tools()
    xtools = []
    for t in tools:
        if t.name == 'CreateTerminalSession':
            xtools.append(t)
    xtools.append(get_daily_string)

    print(xtools)
    prompt = PromptTemplate.from_template("""
    
    {system}
    
    You have access to the following tools:

    {tools}

    Use the following format:

    Question: the question you must answer
    If you want to use tools:
        Thought: always reason what to do
        Action: the action to take, must be one of [{tool_names}]
        Action Input: a **JSON object** with named arguments required by the tool
        Observation: the result of the action
    If no tool is needed:
        Thought: what you are thinking
    ... (this Thought/Action/... can repeat N times)
    Final Answer: the final answer to the original question

    Here is the conversation history:
    {chat_history}
    
    User message: {user}
    {agent_scratchpad}
    """)
    op = MCPReactParser()
    agent = create_react_agent(model_openai, xtools, prompt, output_parser=op)
    message_history = MongoDBChatMessageHistory(
        connection_string=MONGO_CONNECTION_STRING,
        session_id=USER_SESSION_ID,
        database_name=MONGO_DB_NAME,
        collection_name=MONGO_COLLECTION_NAME
    )
    mongodb_memory = ConversationBufferWindowMemory(
        chat_memory=message_history,
        memory_key=MEMORY_KEY,
        k=10,
        return_messages=False,
    )
    agent_executor = AgentExecutor(
        agent=agent,
        tools=xtools,
        memory=mongodb_memory,
        handle_parsing_errors=True,
        verbose=True,

    )
    response = await agent_executor.ainvoke({
        'system': 'you are a helpful assistant',
        'user': 'hello, please create a terminal session with label \'x\' and show me the session id',
    })
    print(response)
    r2 = await agent_executor.ainvoke({
        'system': 'you are a helpful assistant',
        'user': 'hello, please show me today\'s daily string'
    })
    print(r2)
    #print(s1)
    #print(s2)
    #print(s3)
    #print(s4)

if __name__ == '__main__':
    asyncio.run(main())