import os
import warnings
import time
import pickle
import random
import numpy as np

# Zittiamo i warning
warnings.filterwarnings("ignore")

from dotenv import load_dotenv
from typing import TypedDict, Literal, Annotated
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langgraph.graph.message import add_messages
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_google_genai import ChatGoogleGenerativeAI
from pydantic import BaseModel, Field
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver

load_dotenv()

print("--- Inizializzazione Sistema ---")

# --- 1. SETUP MOTORI DI RICERCA ---
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

# A) Caricamento Database PRODOTTI
try:
    product_store = FAISS.load_local("faiss_index_fashion", embeddings, allow_dangerous_deserialization=True)
    # Aumentiamo K a 20 per avere più candidati
    product_retriever = product_store.as_retriever(search_kwargs={"k": 20})
    print("✅ Catalogo Prodotti: CARICATO")
except Exception as e:
    print(f"❌ Catalogo Prodotti: NON TROVATO ({e})")
    product_retriever = None

# B) Caricamento Database ROUTER
try:
    router_store = FAISS.load_local("faiss_router_index", embeddings, allow_dangerous_deserialization=True)
    router_retriever = router_store.as_retriever(search_kwargs={"k": 5})
    print("✅ Cervello Router: CARICATO")
except Exception as e:
    print(f"❌ Cervello Router: NON TROVATO ({e})")
    router_retriever = None

# C) Caricamento Recommender LightFM
rec_model = None
rec_idx_map = {}
try:
    with open("recommender_model.pkl", "rb") as f:
        rec_package = pickle.load(f)
        rec_model = rec_package["model"]
        if "item_id_map" in rec_package: 
             rec_item_map = rec_package["item_id_map"]
             rec_idx_map = rec_package["idx_to_item"]
        else:
             rec_idx_map = {}
    print("✅ Recommender Engine: CARICATO (LightFM)")
except Exception as e:
    print(f"❌ Recommender Engine: NON TROVATO ({e})")


# --- 2. SETUP LLM (Gemini) ---
llm = ChatGoogleGenerativeAI(model="gemini-exp-1206", temperature=0)

# --- 3. DEFINIZIONE STRUTTURE ---

class RouteDecision(BaseModel):
    step: Literal["agent_saluti", "agent_ricerca", "agent_recommender", "agent_default"] = Field(
        description="L'agente a cui assegnare la richiesta"
    )
    reasoning: str = Field(description="Motivo della scelta")

class AgentState(TypedDict):
    messages: Annotated[list, add_messages]
    decision: str

# Funzione ausiliaria per evitare Rate Limit (Pause)
def safe_sleep():
    time.sleep(1.5) # Pausa di 1.5 secondi tra le chiamate

# --- 4. I NODI (GLI AGENTI) ---

def nodo_router(state: AgentState):
    print(f"\n--- [ROUTER] Analisi semantica in corso... ---")
    safe_sleep() # Freno
    last_msg = state["messages"][-1].content
    
    examples_context = ""
    if router_retriever:
        try:
            similar_docs = router_retriever.invoke(last_msg)
            ex_list = []
            for d in similar_docs:
                ex_list.append(f"- '{d.page_content}' -> {d.metadata['target_agent']}")
            examples_context = "\n".join(ex_list)
        except: pass

    sys_msg = SystemMessage(content=f"""Sei il Router. Scegli l'agente:
    1. agent_ricerca: Per prodotti specifici (scarpe rosse, jeans, taglia M).
    2. agent_recommender: Per CONSIGLI generici, idee regalo, outfit.
    3. agent_default: Assistenza, spedizioni, resi.
    4. agent_saluti: Ciao, grazie.

    ESEMPI:
    {examples_context}
    """)
    
    try:
        router_llm = llm.with_structured_output(RouteDecision)
        response = router_llm.invoke([sys_msg] + state["messages"])
        print(f"   Query Utente: '{last_msg}'")
        print(f"   Decisione AI: {response.step.upper()}")
        return {"decision": response.step}
    except Exception as e:
        print(f"⚠️ Errore Router ({e}), vado su Default.")
        return {"decision": "agent_default"}

def nodo_ricerca(state: AgentState):
    print("--- [AGENT RICERCA] 1. Traduzione Query... ---")
    safe_sleep() # Freno
    
    if not product_retriever:
        return {"messages": [AIMessage(content="Errore: Il catalogo prodotti non è disponibile.")]}
    
    user_query = state["messages"][-1].content
    
    # TRADUZIONE
    try:
        translation_prompt = f"""Task: Translate search query to English keywords for Amazon.
        Input: "{user_query}"
        Output keywords only:"""
        english_query_msg = llm.invoke(translation_prompt)
        english_query = english_query_msg.content.strip()
        print(f"   Originale: '{user_query}' -> Ricerca interna: '{english_query}'")
    except:
        english_query = user_query # Fallback
    
    # RICERCA
    print("--- [AGENT RICERCA] 2. Interrogazione Catalogo... ---")
    results = product_retriever.invoke(english_query)
    
    products_text = ""
    for doc in results:
        meta = doc.metadata
        img_info = "Sì" if meta.get('image_url') else "No"
        products_text += f"- Title: {meta.get('title')} | Price: {meta.get('price')} | Image: {img_info}\n"
    
    # RISPOSTA
    print("--- [AGENT RICERCA] 3. Generazione Risposta... ---")
    safe_sleep() # Freno
    
    prompt = f"""Sei un personal shopper.
    Query Utente (IT): "{user_query}"
    Prodotti Trovati (EN):
    {products_text}
    
    1. Seleziona SOLO i prodotti pertinenti alla richiesta italiana.
    2. Rispondi in Italiano descrivendoli.
    3. Se non trovi nulla, dillo.
    """
    
    response = llm.invoke(prompt)
    return {"messages": [response]}

def nodo_recommender(state: AgentState):
    print("--- [AGENT RECOMMENDER] Motore LightFM Attivo... ---")
    
    if not rec_model:
        return {"messages": [AIMessage(content="Il sistema di raccomandazione è offline.")]}
    
    try:
        # Simulazione Item-to-Item
        random_internal_id = random.choice(list(rec_idx_map.keys()))
        anchor_asin = rec_idx_map[random_internal_id]
        
        item_vectors = rec_model.item_embeddings
        target_vector = item_vectors[random_internal_id]
        scores = item_vectors.dot(target_vector)
        top_indices = np.argsort(-scores)[1:4]
        
        suggestions_text = ""
        for idx in top_indices:
            asin = rec_idx_map.get(idx, "N/A")
            suggestions_text += f"- Prodotto ID: {asin}\n"
            
        msg = f"""Ho analizzato i tuoi gusti (Collaborative Filtering LightFM). 🧠
        Ecco alcuni prodotti consigliati basati sui trend attuali:
        
        {suggestions_text}"""
        
    except Exception as e:
        msg = f"Errore nel calcolo dei consigli: {e}"

    return {"messages": [AIMessage(content=msg)]}

def nodo_saluti(state):
    return {"messages": [AIMessage(content="Ciao! Sono il tuo assistente fashion. Posso cercare prodotti o darti consigli.")]}

def nodo_default(state):
    return {"messages": [AIMessage(content="Per assistenza su ordini o resi, contatta il supporto umano.")]}

# --- 5. COSTRUZIONE GRAFO ---
workflow = StateGraph(AgentState)

workflow.add_node("router", nodo_router)
workflow.add_node("agent_saluti", nodo_saluti)
workflow.add_node("agent_ricerca", nodo_ricerca)
workflow.add_node("agent_recommender", nodo_recommender)
workflow.add_node("agent_default", nodo_default)

workflow.set_entry_point("router")
workflow.add_conditional_edges("router", lambda x: x["decision"])
workflow.add_edge("agent_saluti", END)
workflow.add_edge("agent_ricerca", END)
workflow.add_edge("agent_recommender", END)
workflow.add_edge("agent_default", END)

memory = MemorySaver()
app = workflow.compile(checkpointer=memory)

# --- 6. LOOP INTERATTIVO ---
config = {"configurable": {"thread_id": "demo_final_v5"}}

print("\n--- SISTEMA PRONTO ---")
while True:
    try:
        user_input = input("Tu: ")
        if user_input.lower() in ["q", "esci", "exit"]: break
        
        result = app.invoke({"messages": [HumanMessage(content=user_input)]}, config=config)
        print(f"Bot: {result['messages'][-1].content}\n")
    except Exception as e:
        print(f"ERRORE DI SISTEMA: {e}")