from typing import TypedDict, Literal
from langgraph.graph import StateGraph, END

# 1. LO STATO
class RouterState(TypedDict):
    text: str
    classification: str # Qui salviamo la decisione del router ('saluto' o 'ricerca')

# 2. I NODI
def nodo_router(state: RouterState):
    print(f"--- [ROUTER] Analizzo: '{state['text']}' ---")
    
    # Logica simulata (nella tesi qui ci sarà l'LLM)
    msg = state['text'].lower()
    if "scarpe" in msg or "prodotto" in msg:
        decisione = "ricerca"
    elif "ciao" in msg or "buongiorno" in msg:
        decisione = "saluto"
    else:
        decisione = "generico"
        
    print(f"   Deciso: vado su -> {decisione}")
    # Restituisco la classificazione aggiornando lo stato
    return {"classification": decisione}

def nodo_saluti(state: RouterState):
    return {"text": "Ciao! Sono il tuo assistente virtuale. Come posso aiutarti oggi?"}

def nodo_ricerca(state: RouterState):
    return {"text": "Sto cercando nel catalogo i prodotti richiesti... 🔍"}

def nodo_default(state: RouterState):
    return {"text": "Non ho capito, puoi ripetere?"}

# 3. IL GRAFO
workflow = StateGraph(RouterState)

# Aggiungiamo i nodi
workflow.add_node("router", nodo_router)
workflow.add_node("agent_saluti", nodo_saluti)
workflow.add_node("agent_ricerca", nodo_ricerca)
workflow.add_node("agent_default", nodo_default)

# Entry point
workflow.set_entry_point("router")

# 4. EDGES CONDIZIONALI (La parte magica)
# Funzione che legge lo stato e restituisce il NOME del nodo successivo
def decidi_percorso(state):
    if state["classification"] == "saluto":
        return "agent_saluti"
    elif state["classification"] == "ricerca":
        return "agent_ricerca"
    else:
        return "agent_default"

# Aggiungiamo la logica condizionale al router
workflow.add_conditional_edges(
    "router",          # Da dove parte la decisione?
    decidi_percorso    # Quale funzione decide?
    # (Opzionale: mapping esplicito per chiarezza, qui lo ometto perché i nomi coincidono o li gestisce la funzione)
)

# Chiudiamo i flussi (tutti vanno a END per ora)
workflow.add_edge("agent_saluti", END)
workflow.add_edge("agent_ricerca", END)
workflow.add_edge("agent_default", END)

# 5. ESECUZIONE
app = workflow.compile()

# Test 1: Saluto
print("\n--- TEST 1 ---")
print(app.invoke({"text": "Buongiorno, ci sei?"})['text'])

# Test 2: Ricerca
print("\n--- TEST 2 ---")
print(app.invoke({"text": "Vorrei delle scarpe rosse"})['text'])