import asyncio
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from langchain_mcp_adapters.tools import load_mcp_tools
from llm.models import get_chat_model
from langchain.agents import create_agent
from langchain.agents.middleware import wrap_tool_call
from langchain.messages import ToolMessage
from langgraph.checkpoint.memory import InMemorySaver
from data_models.schemas import PDFResponseFormat, Context, InvoiceResponseFormat

@wrap_tool_call
async def handle_tool_errors(request, handler):
    """Handle tool execution errors with custom messages."""
    try:
        return await handler(request)
    except Exception as e:
        # print(f"\n[DEBUG] TOOL ERROR: {type(e).__name__}: {str(e)}\n", flush=True)
        # Return a custom error message to the model
        return ToolMessage(
            content=f"Tool error: Please check your input and try again. ({str(e)})",
            tool_call_id=request.tool_call["id"]
        )

async def run_agent():
    llm = await get_chat_model("gemini")

    import sys
    server_params = StdioServerParameters(
        command=sys.executable,
        args=["server.py"], 
        env=None
    )

    print("Connecting to MCP Server...", flush=True)
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            # Initialize the handshake
            await session.initialize()
            
            # 4. Magically fetch all tools from the server!
            mcp_tools = await load_mcp_tools(session)
            print(f"Loaded tools from server: {[tool.name for tool in mcp_tools]}", flush=True)

            checkpointer = InMemorySaver()
            config = {"configurable": {"thread_id": "1"}}

            agent = create_agent(
                model=llm,
                tools=mcp_tools,
                checkpointer=checkpointer,
                system_prompt=(
                    "You are a professional data extraction assistant. And expert in analysis data from inventory, invoices, purchase order, and shipping order "
                    "You have 2 main roles "
                    "Role 1: Data Extractor when user asking for extract and save data "
                    "Your workflow is: "
                    "1. Searching incoming files in the folder given by user using search_files tool "
                    "2. Load and extract text using read_file_content tool "
                    "3. Extract data from the text using invoice_structured_output or inventory_structured_output tool "
                    "4. Export data to tabular format using tabular_parser tool "
                    "5. return the final answer "
                    "Role 2: Data Analyzer when user asking for analysis data "
                    "Your workflow is: "
                    "1. Searching related files using search_files tool in the output folder "
                    "2. Load and extract text using read_file_content tool "
                    "3. Analyze data using analyze_inventory_health, analyze_shipping_efficiency, analyze_sales_performance tools upon related user query "
                    "4. return the final answer "
                    "You can also chat with user if user asking for something else "
                ),
                middleware=[handle_tool_errors],
                name="Sales Data Agent"
            )

            print("\nAgent is ready! (Type 'quit' to exit)", flush=True)
            while True:
                user_input = input("\nYou: ")
                if user_input.lower() in ["quit", "exit"]:
                    break

                messages = {"messages": [{"role": "user", "content": user_input}]}
                is_thinking = False

                print("\n" + "="*20, flush=True)
                async for event in agent.astream_events(
                    messages, 
                    config=config, 
                    version="v2",
                    context=Context(user_id="1")
                ):
                    kind = event["event"]

                    if kind == "on_chat_model_stream":
                        content = event["data"]["chunk"].content
                        metadata = event["data"]["chunk"].additional_kwargs
                        
                        if "reasoning_content" in metadata:
                            print(f"{metadata['reasoning_content']}", end="", flush=True)
                        
                        elif isinstance(content, str):
                            if "" in content:
                                is_thinking = True
                                print("[Thinking]: ", end="", flush=True)
                                content = content.replace("", "")
                            
                            if "" in content:
                                is_thinking = False
                                content = content.replace("", "")
                                print(f"{content}\n", end="", flush=True)
                                continue

                            if is_thinking:
                                print(f"{content}", end="", flush=True)
                            else:
                                print(content, end="", flush=True)

                    elif kind == "on_chat_model_end":
                        output = event["data"]["output"]
                        if hasattr(output, 'tool_calls') and output.tool_calls:
                            for tool_call in output.tool_calls:
                                print(f"\nTool Call: {tool_call['name']}", flush=True)
                                print(f"Args: {tool_call['args']}", flush=True)
                    
                    elif kind == "on_tool_end":
                        print(f"\nTool Result: {event['data'].get('output')}", flush=True)
                        
                final_state = await agent.aget_state(config)
                if "messages" in final_state.values:
                    last_msg = final_state.values["messages"][-1]
                    if hasattr(last_msg, 'tool_calls') and last_msg.tool_calls:
                        for tool_call in last_msg.tool_calls:
                            if tool_call['name'] in [PDFResponseFormat.__name__, InvoiceResponseFormat.__name__, "structured_output", "final_answer"]:
                                print("\nFinal Structured Output:", flush=True)
                                print(f"Data: {tool_call['args']}", flush=True)

if __name__ == "__main__":
    asyncio.run(run_agent())