from documents.services.document_loader import load_pdf, search_pdf_files
from documents.services.parser_output import tabular_parser
from llm.schemas import PDFResponseFormat, Context, InvoiceResponseFormat
from llm.models.llm_models import get_chat_model
from langgraph.checkpoint.memory import InMemorySaver
from langchain.agents import create_agent
from langchain.agents.structured_output import ToolStrategy
from langchain.agents.middleware import wrap_tool_call
from langchain.messages import ToolMessage
import asyncio

@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:
        # 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 main():
    tools = [load_pdf, search_pdf_files, tabular_parser]
    llm = await get_chat_model("local")
    checkpointer = InMemorySaver()
    config = {"configurable": {"thread_id": "1"}}

    agent = create_agent(
        model=llm,
        tools=tools,
        checkpointer=checkpointer,
        response_format=ToolStrategy(InvoiceResponseFormat),
        system_prompt=(
            "You are a professional PDF data extraction assistant. "
            "Your workflow is: "
            "1. Search for the requested PDF file using search_pdf_files. "
            "2. Load and extract text using load_pdf. "
            "3. Extract all required invoice fields (Order ID, Customer ID, Products, etc.). "
            "4. Call tabular_parser to export the extracted data to a file (CSV or JSON as requested). "
            "5. FINAL STEP: Use the provided structured output format (InvoiceResponseFormat) to return the full extracted data as your final response. "
            "CRITICAL: Always use tabular_parser BEFORE providing your final response format."
        ),
        middleware=[handle_tool_errors],
        name="PDF Data Extractor"
    )

    print("\n🤖 Agent is ready! (Type 'quit' to exit)")

    while True:
        user_input = input("\nYou: ")
        if user_input.lower() in ["quit", "exit"]:
            break

        messages = {"messages": [{"role": "user", "content": user_input}]}
        # Tracking state to handle  tags if they arrive as raw text
        is_thinking = False

        print("\n" + "="*20)
        # Use astream_events (version 2) for granular control over streaming
        async for event in agent.astream_events(
            messages, 
            config=config, 
            version="v2",
            context=Context(user_id="1")
        ):
            kind = event["event"]

            # 1. SHOW MODEL THINKING
            # Most 2026 local models (like DeepSeek or Gemini) stream reasoning here
            if kind == "on_chat_model_stream":
                content = event["data"]["chunk"].content
                metadata = event["data"]["chunk"].additional_kwargs
                
                # Check for reasoning/thinking tokens (depends on your local provider)
                if "reasoning_content" in metadata:
                    print(f"{metadata['reasoning_content']}", end="", flush=True)
                
                # B. Handle raw  tags (common if the local server doesn't parse them)
                elif isinstance(content, str):
                    if "" in content:
                        is_thinking = True
                        # Print everything after the tag in grey
                        print("[Thinking]: ", end="", flush=True)
                        content = content.replace("", "")
                    
                    if "" in content:
                        is_thinking = False
                        # Close the grey color and move to a new line
                        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)

            # 3. SHOW TOOL CALLS (As they are decided)
            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"Tool Call: {tool_call['name']}")
                        print(f"Args: {tool_call['args']}")
            
            # 4. SHOW TOOL RESULTS
            elif kind == "on_tool_end":
                print(f"Tool Result: {event['data'].get('output')}")
        # 5. FETCH FINAL STRUCTURED RESULT
        # After streaming is done, fetch the last message to see if it's the structured output
        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("\n✅ Final Structured Output:")
                        print(f"Data: {tool_call['args']}")
            elif hasattr(last_msg, 'content') and last_msg.content:
                 # Check if the content itself is a JSON of the format (fallback)
                 print("\n⚠️ Note: Received text response instead of tool call. Check if format is respected.")

        print("\n" + "="*20)

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