from typing import Annotated, List, TypedDict
from langchain_google_genai import ChatGoogleGenerativeAI
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import ToolNode
from src.tools import tools

# The Local Knowledge Registry (Update this whenever you add new data types)
# Manual as of now
LOCAL_MANIFEST = {
    "topics": ["HR Policies", "Project X Design Docs", "Q3 Financials", "Employee Handbook"],
    "date_range": "Documents updated as of Dec 2024",
    "domain": "Internal Corporate Knowledge"
}

SYSTEM_PROMPT = f"""
You are an expert Research Assistant. You have access to:
1. INTERNAL DATA: {LOCAL_MANIFEST['topics']}. (Use 'local_research_tool')
2. EXTERNAL DATA: The entire internet via duckduckgosearch. (Use 'web_search_tool')

GUIDELINES:
- Given the user's technical question and the fact that our internal documents are insufficient, generate a generic search query for the internet that does NOT include any proprietary names or internal details.
- If a query is about {LOCAL_MANIFEST['topics']}, try LOCAL first.
- If a query is TECHNICAL (e.g., PyTorch, Python APIs) or REAL-TIME, go to WEB immediately.
- If the query is ambiguous, try LOCAL first, then fallback to WEB if the results are empty or low confidence.
"""


class AgentState(TypedDict):
    """MNC Agent state with intent tracking and sufficiency grading."""
    messages: Annotated[List, "Chat history"]
    intent: str
    is_sufficient: bool

# Brain: Gemini 2.0 Flash (high-speed agentic reasoning)
llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
llm_with_tools = llm.bind_tools(tools)

def router(state: AgentState):
    """Classifies user intent to prioritize retrieval paths."""
    query = state['messages'][-1].content
    prompt = f"Categorize intent: TECHNICAL (API/Docs), INTERNAL (Proprietary), or REALTIME. Query: {query}"
    response = llm.invoke(prompt)
    intent = "TECHNICAL" if any(x in response.content.upper() for x in ["TECHNICAL", "REALTIME"]) else "INTERNAL"
    return {"intent": intent}

def call_model(state: AgentState):
    """Invokes Gemini with tools based on intent and history."""
    return {"messages": [llm_with_tools.invoke(state['messages'])]}

# Orchestration Graph
workflow = StateGraph(AgentState)

workflow.add_node("router", router)
workflow.add_node("llm", call_model)
workflow.add_node("tools", ToolNode(tools))

workflow.add_edge(START, "router")
workflow.add_edge("router", "llm")

# Self-Correction Loop
def should_continue(state: AgentState):
    last_msg = state["messages"][-1]
    return "tools" if last_msg.tool_calls else END

workflow.add_conditional_edges("llm", should_continue)
workflow.add_edge("tools", "llm")

app = workflow.compile()