"""
Builds the FAISS index for the Router Agent.

Data sources (combined into a single index):
1. HuggingFace dataset 'vansh-khaneja/ecommerce-intent-routing' (~2750 examples)
   - Generic e-commerce (laptop, smartphone, etc.)
2. Fashion sentences generated with GPT-4o-mini (artifacts/fashion_examples.json, ~500 examples)
   - Specific to the clothing/shoes/accessories domain

5 original classes: product, sales, support, account, bugs.
The classes are kept as-is in the FAISS index (in the 'intent_class' field).
The intent → agent mapping is done at runtime in the Router Agent, not here.
This avoids artificial imbalance (e.g. merging 3 classes into 'default' = 60% of the dataset).

It also adds a 'recommender_hint' field (True/False) for the product/sales classes:
a keyword heuristic that identifies recommendation queries (recommend, suggest, best, gift...).
The Router can use this hint to distinguish search from recommendation.
"""

import json
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

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

RECOMMENDER_KEYWORDS = [
    "recommend", "suggest", "best", "top rated", "ideas for",
    "what should i", "good for", "advice", "opinion", "popular",
    "trending", "look for", "matching", "outfit", "gift", "review"
]


def has_recommender_intent(text):
    """Heuristic: looks for keywords that indicate a request for advice/opinion."""
    text_lower = text.lower()
    return any(kw in text_lower for kw in RECOMMENDER_KEYWORDS)


def ingest_router_data():
    print(f"--- Downloading dataset '{DATASET_NAME}'... ---")
    try:
        dataset = load_dataset(DATASET_NAME, split="train")
        print(f"Columns found: {dataset.column_names}")
    except Exception as e:
        print(f"Dataset download error: {e}")
        return

    print("--- Loading Local Embeddings model (HuggingFace)... ---")
    embeddings = HuggingFaceEmbeddings(model_name="paraphrase-multilingual-MiniLM-L12-v2")

    documents = []
    counts = {}

    print(f"--- Starting indexing... ---")

    for row in dataset:
        text = row['query']
        raw_label = row['route']

        # Converts numeric label to string
        label_text = str(raw_label)
        if isinstance(raw_label, int):
            try:
                features = dataset.features['route']
                label_text = features.int2str(raw_label)
            except Exception:
                pass

        intent_class = label_text.lower().strip()

        # Count per original class
        counts[intent_class] = counts.get(intent_class, 0) + 1

        # For product/sales, hint whether it is a recommendation query
        recommender_hint = False
        if intent_class in ("product", "sales"):
            recommender_hint = has_recommender_intent(text)

        doc = Document(
            page_content=text,
            metadata={
                "intent_class": intent_class,
                "recommender_hint": recommender_hint,
            }
        )
        documents.append(doc)

    print(f"--- HuggingFace Dataset Statistics ---")
    print(f"HuggingFace examples: {len(documents)}")
    for cls, count in sorted(counts.items(), key=lambda x: -x[1]):
        pct = count / len(documents) * 100
        print(f"  {cls}: {count} ({pct:.1f}%)")

    # --- Load additional fashion sentences ---
    if os.path.exists(FASHION_EXAMPLES_PATH):
        print(f"\n--- Loading fashion sentences from {FASHION_EXAMPLES_PATH}... ---")
        with open(FASHION_EXAMPLES_PATH, "r", encoding="utf-8") as f:
            fashion_data = json.load(f)

        fashion_counts = {}
        for item in fashion_data:
            text = item["query"]
            intent_class = item["route"].lower().strip()

            fashion_counts[intent_class] = fashion_counts.get(intent_class, 0) + 1
            counts[intent_class] = counts.get(intent_class, 0) + 1

            recommender_hint = False
            if intent_class in ("product", "sales"):
                recommender_hint = has_recommender_intent(text)

            doc = Document(
                page_content=text,
                metadata={
                    "intent_class": intent_class,
                    "recommender_hint": recommender_hint,
                }
            )
            documents.append(doc)

        print(f"Fashion sentences added: {sum(fashion_counts.values())}")
        for cls, count in sorted(fashion_counts.items()):
            print(f"  {cls}: {count}")
    else:
        print(f"\n--- No fashion_examples.json file found, only HuggingFace dataset ---")

    print(f"\n--- Total Statistics ---")
    print(f"Total examples: {len(documents)}")
    for cls, count in sorted(counts.items(), key=lambda x: -x[1]):
        pct = count / len(documents) * 100
        print(f"  {cls}: {count} ({pct:.1f}%)")

    print(f"--- Saving FAISS index... ---")
    vector_store = FAISS.from_documents(documents, embeddings)
    vector_store.save_local(ROUTER_INDEX_FOLDER)
    print(f"--- DONE! ---")


if __name__ == "__main__":
    ingest_router_data()