"""
LangGraph graph definition — the heart of the multi-agent orchestration.

This module:
1. Loads all artifacts from disk (FAISS, LightFM, metadata, DinoV2)
2. Creates the agent nodes by injecting the dependencies (factory pattern)
3. Defines the graph with conditional routing:

    [User] → Router → agent_search      → [Response]
                     → agent_recommender  → [Response]
                     → agent_visual       → [Response]
                     → agent_greeting       → [Response]
                     → agent_default      → [Response]

The Router decides the destination. Each agent responds and the flow ends (END).
The MemorySaver keeps the message history across conversation turns.

Exports: `app` (compiled graph), imported by main.py.
"""

import warnings
warnings.filterwarnings("ignore")

from typing import TypedDict, Annotated
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langgraph.checkpoint.memory import MemorySaver

from utils.llm import llm
from utils.loaders import (
    load_product_retriever, load_router_retriever, load_recommender,
    load_metadata_lookup, load_visual_index
)
from agents import (
    create_router_node, create_search_node, create_recommender_node,
    create_visual_node, create_greeting_node, default_node
)


# --- Definition of the state shared across nodes ---
class AgentState(TypedDict):
    messages: Annotated[list, add_messages]  # Message history (accumulative)
    decision: str                             # Destination chosen by the Router
    resolved_query: str                       # Query rewritten by the Router with conversational context resolved
    image_path: str                           # Image path for the Visual Agent (if present)
    products: list                            # List of structured products [{asin, title, brand, price, similarity}]
    greeting_type: str                        # Greeting type: greeting, thanks, goodbye
    anchor: dict                              # Recommender anchor {asin, title} to display in the frontend
    recommend_asin: str                       # ASIN selected by the user for recommendation (card click)


# --- Loading artifacts from disk ---
print("--- Initializing System ---")
product_retriever = load_product_retriever()   # FAISS product catalog
router_retriever = load_router_retriever()     # FAISS router few-shot examples
rec_model, rec_idx_map, rec_item_features = load_recommender()  # LightFM + mappings + features
metadata_lookup = load_metadata_lookup()       # ASIN → product details
vis_model, vis_processor, vis_index, vis_asin_map, vis_device = load_visual_index()  # DinoV2 + visual FAISS

# --- Creating agent nodes (dependency injection) ---
router_node = create_router_node(llm, router_retriever)
search_node = create_search_node(product_retriever)
recommender_node = create_recommender_node(rec_model, rec_idx_map, metadata_lookup, rec_item_features)
visual_node = create_visual_node(vis_model, vis_processor, vis_index, vis_asin_map, metadata_lookup, vis_device)
greeting_node = create_greeting_node()

# --- Building the graph ---
workflow = StateGraph(AgentState)

# Registering nodes
workflow.add_node("router", router_node)
workflow.add_node("agent_greeting", greeting_node)
workflow.add_node("agent_search", search_node)
workflow.add_node("agent_recommender", recommender_node)
workflow.add_node("agent_visual", visual_node)
workflow.add_node("agent_default", default_node)

# The Router is the entry point
workflow.set_entry_point("router")

# Conditional routing: the Router writes "decision" into the state,
# LangGraph uses that value to choose the next node
workflow.add_conditional_edges("router", lambda x: x["decision"])

# All agents end the flow after the response
workflow.add_edge("agent_greeting", END)
workflow.add_edge("agent_search", END)
workflow.add_edge("agent_recommender", END)
workflow.add_edge("agent_visual", END)
workflow.add_edge("agent_default", END)

# Compile the graph with the checkpointer for session memory
memory = MemorySaver()
app = workflow.compile(checkpointer=memory)