import os
import pandas as pd
from langchain_core.messages import HumanMessage, AIMessage
from langgraph.graph import StateGraph, MessagesState
from langchain_community.llms import HuggingFaceEndpoint  
from tools import TOOLS  

# --- Read local QA data for retriever ---
QA_PATH = "metadata.jsonl"
qa_pairs = pd.read_json(QA_PATH, lines=True)
qa_dict = {
    row["Question"].strip(): row["Final answer"].strip()
    for _, row in qa_pairs.iterrows()
}

def build_graph():
    # Initialize Mistral model - CORRECTED CONFIGURATION
    llm = HuggingFaceEndpoint(
        repo_id="mistralai/Mistral-7B-Instruct-v0.3",
        task="text-generation",
        huggingfacehub_api_token=os.environ["HF_TOKEN"],
        # model_kwargs={
        #     "max_new_tokens": 512,
        #     "temperature": 0.1
        # }
    )

    # Retriever node (unchanged)
    def retriever_node(state: MessagesState):
        query = state["messages"][-1].content.strip()
        if query in qa_dict:
            print("✅ Exact match found in retriever.")
            return {"messages": [AIMessage(content=qa_dict[query])]}
        print("🔍 No match. Sending to LLM.")
        return {"messages": state["messages"]}

    # Assistant node (LLM) - UPDATED FOR TEXT-GENERATION
    def assistant_node(state: MessagesState):
        query = state["messages"][-1].content.strip()

        # Format system prompt for text-generation
        system_prompt = (
            "You are a helpful assistant evaluated by the GAIA benchmark. "
            "Only return the final answer, with no explanations. "
            "- No prefixes like 'Final answer:' "
            "- If it's a list, output comma-separated "
            "- If unknown, say 'Unknown' "
            "- Never justify or explain"
        )
        
        # Format prompt for text-generation model
        prompt = f"<s>[INST] {system_prompt}\n\n{query} [/INST]"
        response = llm.invoke(prompt).strip()

        # Clean up response
        for tag in ("Final answer:", "Answer:", "assistant:"):
            if response.lower().startswith(tag.lower()):
                response = response[len(tag):].strip()
        
        return {"messages": [AIMessage(content=response.strip())]}

    # Tool node (unchanged)
    def tool_node(state: MessagesState):
        try:
            tool_signal = state.get("tool_call", "")
            _, tool_name, tool_input = tool_signal.split(":", 2)
            tool_name = tool_name.strip()
            tool_input = tool_input.strip()
            tool_fn = TOOLS.get(tool_name)

            if not tool_fn:
                print(f"❌ Unknown tool: {tool_name}")
                return {"messages": [AIMessage(content="Unknown")]}
            
            print(f"🔧 Using tool: {tool_name} with input: {tool_input}")
            tool_result = tool_fn(tool_input)
            return {"messages": [AIMessage(content=str(tool_result))]}

        except Exception as e:
            print(f"⚠️ Tool error: {e}")
            return {"messages": [AIMessage(content="Unknown")]}  # fail-safe

    # Build LangGraph (unchanged)
    builder = StateGraph(MessagesState)
    builder.add_node("retriever", retriever_node)
    builder.add_node("assistant", assistant_node)
    builder.add_node("tool", tool_node)

    builder.set_entry_point("retriever")
    builder.add_edge("retriever", "assistant")
    builder.add_edge("assistant", "tool")
    builder.add_edge("tool", "assistant")
    builder.set_finish_point("assistant")

    return builder.compile()

# Agent class (unchanged)
class BasicAgent:
    def __init__(self):
        print("✅ BasicAgent initialized with retriever + LLM + tools")
        self.graph = build_graph()

    def __call__(self, question: str) -> str:
        print(f"📥 Question: {question[:100]}")
        result = self.graph.invoke({"messages": [HumanMessage(content=question)]})
        answer = result["messages"][-1].content.strip()
        print(f"📤 Answer: {answer}")
        return answer