"""
Loading artifacts from disk.

Each function loads a specific artifact from the artifacts/ folder:
- load_product_retriever()  → FAISS product catalog index
- load_router_retriever()   → FAISS few-shot examples index for the router
- load_recommender()        → LightFM model + mappings + features matrix
- load_metadata_lookup()    → ASIN → product details dictionary
- load_visual_index()       → visual FAISS index + DinoV2 model

All functions handle the case where the artifact does not exist,
returning None or an empty dictionary. This lets the system
start up even with missing artifacts (degraded functionality).
"""

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

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from utils.config import (
    EMBEDDING_MODEL_NAME, FAISS_PRODUCT_INDEX, FAISS_ROUTER_INDEX,
    RECOMMENDER_MODEL_PATH, METADATA_PATH, PRODUCT_SEARCH_K, ROUTER_SEARCH_K,
    FAISS_VISUAL_INDEX, VISUAL_MODEL_NAME
)

# Embedding model shared across all FAISS indexes
embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)


def load_product_retriever():
    """Load the FAISS product catalog index (100k items).
    Returns a LangChain retriever that, given a query, returns the k most similar products."""
    try:
        store = FAISS.load_local(FAISS_PRODUCT_INDEX, embeddings, allow_dangerous_deserialization=True)
        print("  Product Catalog: LOADED")
        return store.as_retriever(search_kwargs={"k": PRODUCT_SEARCH_K})
    except Exception as e:
        print(f"  Product Catalog: NOT FOUND ({e})")
        return None


def load_router_retriever():
    """Load the FAISS few-shot examples index for the router (~2750 examples).
    Returns the direct FAISS store (not the wrapped retriever) to allow
    the Router to use similarity_search_with_score() in distance-weighted kNN voting."""
    try:
        store = FAISS.load_local(FAISS_ROUTER_INDEX, embeddings, allow_dangerous_deserialization=True)
        print("  Router Brain: LOADED")
        return store
    except Exception as e:
        print(f"  Router Brain: NOT FOUND ({e})")
        return None


def load_recommender():
    """Load the hybrid LightFM model and its mappings.

    The pickle contains:
    - model: the trained LightFM model
    - idx_to_item: maps internal index → ASIN
    - item_features_matrix: sparse feature matrix (categories + price),
      required for get_item_representations(). None if the model is not hybrid.

    Returns: (model, idx_to_item, item_features_matrix)
    """
    try:
        with open(RECOMMENDER_MODEL_PATH, "rb") as f:
            package = pickle.load(f)
            model = package["model"]
            idx_map = package.get("idx_to_item", {})
            item_features = package.get("item_features_matrix", None)
        print("  Recommender Engine: LOADED (hybrid LightFM)")
        return model, idx_map, item_features
    except Exception as e:
        print(f"  Recommender Engine: NOT FOUND ({e})")
        return None, {}, None


def load_metadata_lookup():
    """Load the metadata of the 100k products and index it by ASIN.
    Returns a dictionary {ASIN: {title, brand, price, categories, image_url}}.
    Used by the Recommender Agent to display real product names."""
    try:
        with open(METADATA_PATH, "r", encoding="utf-8") as f:
            products = json.load(f)
        lookup = {p["asin"]: p for p in products}
        print(f"  Metadata Catalog: LOADED ({len(lookup)} products)")
        return lookup
    except Exception as e:
        print(f"  Metadata Catalog: NOT FOUND ({e})")
        return {}


def load_visual_index():
    """Load the visual FAISS index and the Trendyol DinoV2 model.

    Returns: (model, processor, faiss_index, asin_map, device)
    - model: DinoV2 model (to compute the query image embedding)
    - processor: image preprocessor
    - faiss_index: FAISS IndexFlatIP index (cosine similarity, 256d)
    - asin_map: list [ASIN] where the index corresponds to the FAISS position
    - device: torch.device used (cpu/mps/cuda)

    If the index or the model are not available, returns tuples of None.
    """
    index_path = os.path.join(FAISS_VISUAL_INDEX, "index.faiss")
    map_path = os.path.join(FAISS_VISUAL_INDEX, "asin_map.json")

    # Check that the index exists before loading the model (heavy)
    if not os.path.exists(index_path) or not os.path.exists(map_path):
        print("  Visual Index: NOT FOUND (run pipeline/7 and 8)")
        return None, None, None, [], None

    try:
        import torch
        import faiss as faiss_lib
        from transformers import AutoModel, AutoImageProcessor

        # Load FAISS index and ASIN map
        faiss_index = faiss_lib.read_index(index_path)
        with open(map_path, "r") as f:
            asin_map = json.load(f)
        print(f"  Visual Index: LOADED ({faiss_index.ntotal} vectors)")

        # Load DinoV2 model
        print(f"  Visual Model: loading {VISUAL_MODEL_NAME}...")
        processor = AutoImageProcessor.from_pretrained(VISUAL_MODEL_NAME, trust_remote_code=True)
        model = AutoModel.from_pretrained(VISUAL_MODEL_NAME, trust_remote_code=True)
        model.eval()

        # Device selection
        if torch.backends.mps.is_available():
            device = torch.device("mps")
        elif torch.cuda.is_available():
            device = torch.device("cuda")
        else:
            device = torch.device("cpu")

        model = model.to(device)
        print(f"  Visual Model: LOADED (device: {device})")

        return model, processor, faiss_index, asin_map, device

    except Exception as e:
        print(f"  Visual Index: ERROR ({e})")
        return None, None, None, [], None