import os
import warnings
warnings.filterwarnings("ignore")

from datasets import load_dataset
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document

# --- CONFIGURAZIONE ---
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
ROUTER_INDEX_FOLDER = os.path.join(PROJECT_ROOT, "artifacts", "faiss_router")
DATASET_NAME = "vansh-khaneja/ecommerce-intent-routing"

def analyze_text_for_recommendation(text):
    """
    Euristica: Cerca keyword che indicano richiesta di consiglio/opinione
    """
    text = text.lower()
    recommender_keywords = [
        "recommend", "suggest", "best", "top rated", "ideas for", 
        "what should i", "good for", "advice", "opinion", "popular",
        "trending", "look for", "matching", "outfit", "gift", "review"
    ]
    
    for kw in recommender_keywords:
        if kw in text:
            return True
    return False

def map_intent_to_agent(intent_label, text_content):
    label_str = str(intent_label).lower()
    
    # 1. CATEGORIA PRODOTTI / VENDITE (Search vs Recommender)
    # Nota: A volte l'etichetta è 'product', a volte 'sales'.
    if "product" in label_str or "sales" in label_str:
        if analyze_text_for_recommendation(text_content):
            return "agent_recommender"
        else:
            return "agent_ricerca"
    
    # 2. CATEGORIE SUPPORTO (Default)
    # account, bugs, support -> agent_default
    elif any(x in label_str for x in ["account", "bug", "support", "order", "track"]):
        return "agent_default"
    
    # Fallback
    return "agent_default"

def ingest_router_data():
    print(f"--- Scaricamento dataset '{DATASET_NAME}'... ---")
    try:
        dataset = load_dataset(DATASET_NAME, split="train")
        print(f"Colonne trovate: {dataset.column_names}")
    except Exception as e:
        print(f"Errore download dataset: {e}")
        return

    print("--- Caricamento modello Embeddings Locale (HuggingFace)... ---")
    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    
    documents = []
    
    # Contatori
    counts = {"agent_ricerca": 0, "agent_recommender": 0, "agent_default": 0}
    
    print(f"--- Inizio indicizzazione e classificazione... ---")
    
    for row in dataset:
        # CORREZIONE QUI: Usiamo i nomi corretti delle colonne
        text = row['query']   # Era 'text'
        raw_label = row['route'] # Era 'label' o 'intent'
        
        # GESTIONE SICURA DELL'ETICHETTA
        # Se l'etichetta è un numero (ID), proviamo a convertirla
        label_text = str(raw_label) 
        if isinstance(raw_label, int):
            try:
                features = dataset.features['route']
                label_text = features.int2str(raw_label)
            except:
                pass
        
        # MAPPING
        target_agent = map_intent_to_agent(label_text, text)
        
        if target_agent in counts: counts[target_agent] += 1
        
        doc = Document(
            page_content=text,
            metadata={
                "target_agent": target_agent,
                "original_category": label_text
            }
        )
        documents.append(doc)
    
    print(f"--- Statistiche Dataset Elaborato ---")
    print(f"Totale esempi: {len(documents)}")
    print(f" -> Assegnati a SEARCH:      {counts['agent_ricerca']}")
    print(f" -> Assegnati a RECOMMENDER: {counts['agent_recommender']}")
    print(f" -> Assegnati a DEFAULT:     {counts['agent_default']}")
    
    print(f"--- Salvataggio indice FAISS... ---")
    vector_store = FAISS.from_documents(documents, embeddings)
    vector_store.save_local(ROUTER_INDEX_FOLDER)
    print(f"--- COMPLETATO! ---")

if __name__ == "__main__":
    ingest_router_data()