"""
Evaluation of the Search Agent — Relevance of the Search Results.

Method: keyword-based relevance evaluation.
For a set of test queries (in English, as the Search Agent receives them from the Router):
1. Runs the FAISS search (same pipeline as the Search Agent)
2. Checks whether the returned products are relevant to the query
   (keywords present in the title OR in the product categories)
3. Computes Precision@k, MRR (Mean Reciprocal Rank), Hit Rate

The FAISS index indexes title + brand + categories (page_content in the builder).
The queries reach the Search Agent already translated into English by the Router.
"""

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

import numpy as np

PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, PROJECT_ROOT)

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from utils.config import EMBEDDING_MODEL_NAME, FAISS_PRODUCT_INDEX, PRODUCT_SEARCH_K

# --- Test queries (in English, as the Search Agent receives them from the Router) ---
# A result is relevant if at least one keyword appears in the title or in the categories.

TEST_QUERIES = [
    # --- Clothing ---
    {"query": "red dress for a party", "keywords": ["dress"]},
    {"query": "men's casual t-shirt", "keywords": ["t-shirt", "tee", "shirt"]},
    {"query": "women's jeans skinny", "keywords": ["jean"]},
    {"query": "winter wool coat", "keywords": ["coat"]},
    {"query": "silk blouse women", "keywords": ["blouse", "shirt", "top"]},
    {"query": "men's hoodie black", "keywords": ["hoodie", "sweatshirt"]},
    {"query": "summer shorts men", "keywords": ["short"]},
    {"query": "yoga pants women", "keywords": ["yoga", "legging", "pant"]},
    {"query": "women's sweater cashmere", "keywords": ["sweater", "pullover"]},
    {"query": "men's formal suit jacket", "keywords": ["suit", "blazer", "jacket"]},
    {"query": "cotton polo shirt men", "keywords": ["polo", "shirt"]},
    {"query": "women's maxi skirt floral", "keywords": ["skirt"]},
    {"query": "men's cargo pants", "keywords": ["pant", "cargo"]},
    {"query": "women's cardigan long sleeve", "keywords": ["cardigan", "sweater"]},
    {"query": "denim jacket women", "keywords": ["jacket", "denim"]},

    # --- Shoes ---
    {"query": "running shoes men", "keywords": ["running", "sneaker", "shoe", "athletic"]},
    {"query": "women's high heels", "keywords": ["heel", "pump"]},
    {"query": "leather boots men", "keywords": ["boot"]},
    {"query": "summer sandals women", "keywords": ["sandal"]},
    {"query": "white sneakers casual", "keywords": ["sneaker"]},
    {"query": "ballet flats women", "keywords": ["flat", "ballet"]},
    {"query": "men's loafers leather", "keywords": ["loafer", "moccasin"]},
    {"query": "winter snow boots", "keywords": ["boot"]},
    {"query": "hiking boots waterproof", "keywords": ["boot", "hiking"]},
    {"query": "women's wedge sandals", "keywords": ["wedge", "sandal"]},

    # --- Accessories ---
    {"query": "sunglasses aviator style", "keywords": ["sunglass", "eyewear"]},
    {"query": "leather handbag women", "keywords": ["handbag", "bag", "purse", "tote"]},
    {"query": "men's baseball cap", "keywords": ["cap", "hat"]},
    {"query": "backpack for school", "keywords": ["backpack"]},
    {"query": "silk scarf women", "keywords": ["scarf"]},
    {"query": "men's leather belt", "keywords": ["belt"]},
    {"query": "women's watch gold", "keywords": ["watch"]},

    # --- Underwear and swimwear ---
    {"query": "women's pajamas cotton", "keywords": ["pajama", "sleepwear", "nightgown", "lounge"]},
    {"query": "women's bikini swimsuit", "keywords": ["bikini", "swimsuit", "swim"]},
    {"query": "sports bra women", "keywords": ["bra", "sports"]},

    # --- More specific queries ---
    {"query": "waterproof rain jacket", "keywords": ["rain", "jacket", "waterproof"]},
    {"query": "baby girl dress pink", "keywords": ["dress", "baby"]},
    {"query": "men's thermal underwear", "keywords": ["thermal", "underwear", "base layer"]},
    {"query": "women's trench coat beige", "keywords": ["trench", "coat"]},
    {"query": "gym shorts men athletic", "keywords": ["short", "athletic", "gym"]},
]


def is_relevant(title, categories, keywords):
    """A product is relevant if at least one keyword appears in the title or in the categories."""
    text = (title + " " + " ".join(categories)).lower()
    return any(kw.lower() in text for kw in keywords)


def evaluate_search(vector_store, metadata_by_asin, k_values):
    """Evaluates the FAISS search over all the test queries."""
    results_all = []

    for tq in TEST_QUERIES:
        query = tq["query"]
        keywords = tq["keywords"]

        max_k = max(k_values)
        docs = vector_store.similarity_search_with_score(query, k=max_k)

        relevance = []
        for doc, score in docs:
            asin = doc.metadata.get("asin", "")
            title = doc.metadata.get("title", "")
            product = metadata_by_asin.get(asin, {})
            cats = product.get("categories", [])

            rel = is_relevant(title, cats, keywords)

            relevance.append({
                "asin": asin,
                "title": title[:70],
                "categories": cats,
                "relevant": rel,
                "score": score,
            })

        results_all.append({
            "query": query,
            "keywords": keywords,
            "results": relevance,
        })

    # --- Metrics computation ---
    metrics = {}
    for k in k_values:
        precisions = []
        for r in results_all:
            top_k = r["results"][:k]
            rel = sum(1 for x in top_k if x["relevant"])
            precisions.append(rel / k)
        metrics[f"precision@{k}"] = np.mean(precisions)

    # MRR
    rrs = []
    for r in results_all:
        rr = 0.0
        for i, x in enumerate(r["results"]):
            if x["relevant"]:
                rr = 1.0 / (i + 1)
                break
        rrs.append(rr)
    metrics["mrr"] = np.mean(rrs)

    # Hit Rate@k
    for k in k_values:
        hits = []
        for r in results_all:
            top_k = r["results"][:k]
            has_hit = any(x["relevant"] for x in top_k)
            hits.append(1.0 if has_hit else 0.0)
        metrics[f"hit_rate@{k}"] = np.mean(hits)

    return metrics, results_all


def main():
    print("=" * 60)
    print("SEARCH AGENT EVALUATION — Result Relevance")
    print(f"Embedding: {EMBEDDING_MODEL_NAME}")
    print(f"Test queries: {len(TEST_QUERIES)} (English)")
    print("=" * 60)

    # Load metadata for the categories
    metadata_path = os.path.join(PROJECT_ROOT, "artifacts", "metadata_cleaned.json")
    print(f"\nLoading metadata...")
    with open(metadata_path, "r", encoding="utf-8") as f:
        metadata_list = json.load(f)
    metadata_by_asin = {p["asin"]: p for p in metadata_list}
    print(f"  {len(metadata_by_asin)} products")

    # Load FAISS index
    print(f"Loading FAISS index...")
    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
    vector_store = FAISS.load_local(
        FAISS_PRODUCT_INDEX, embeddings, allow_dangerous_deserialization=True
    )
    print("  Index loaded")

    k_values = sorted(set([5, 10, PRODUCT_SEARCH_K]))

    print(f"\nEvaluation in progress ({len(TEST_QUERIES)} queries)...")
    metrics, results_all = evaluate_search(vector_store, metadata_by_asin, k_values)

    # --- Report ---
    print(f"\n{'='*60}")
    print(f"AGGREGATE METRICS")
    print(f"{'='*60}")
    print(f"\n{'Metric':<20} {'Value':>10}")
    print(f"{'-'*30}")
    for k in k_values:
        print(f"{'Precision@'+str(k):<20} {metrics[f'precision@{k}']:>10.4f}")
    print(f"{'-'*30}")
    print(f"{'MRR':<20} {metrics['mrr']:>10.4f}")
    print(f"{'-'*30}")
    for k in k_values:
        print(f"{'Hit Rate@'+str(k):<20} {metrics[f'hit_rate@{k}']:>10.4f}")

    # --- Per-query detail ---
    print(f"\n{'='*60}")
    print(f"PER-QUERY DETAIL (top-5)")
    print(f"{'='*60}")

    for r in results_all:
        top5 = r["results"][:5]
        rel_count = sum(1 for x in top5 if x["relevant"])
        precision = rel_count / 5
        status = "OK" if precision >= 0.6 else "!!"
        print(f"\n[{status}] \"{r['query']}\"  (keywords: {r['keywords']})")
        print(f"     Precision@5: {precision:.2f} ({rel_count}/5)")
        for i, x in enumerate(top5):
            mark = "V" if x["relevant"] else "X"
            print(f"       {i+1}. [{mark}] {x['title']}")

    # --- Save to file ---
    results_path = os.path.join(PROJECT_ROOT, "evaluation", "search_results.txt")
    with open(results_path, "w", encoding="utf-8") as f:
        f.write(f"SEARCH AGENT EVALUATION — Result Relevance\n")
        f.write(f"Embedding: {EMBEDDING_MODEL_NAME}\n")
        f.write(f"Test queries: {len(TEST_QUERIES)} (English)\n\n")

        f.write(f"METRICS\n")
        for k in k_values:
            f.write(f"  Precision@{k}: {metrics[f'precision@{k}']:.4f}\n")
        f.write(f"  MRR: {metrics['mrr']:.4f}\n")
        for k in k_values:
            f.write(f"  Hit Rate@{k}: {metrics[f'hit_rate@{k}']:.4f}\n")

        f.write(f"\nPER-QUERY DETAIL\n")
        for r in results_all:
            top_k = r["results"][:PRODUCT_SEARCH_K]
            rel_count = sum(1 for x in top_k if x["relevant"])
            precision = rel_count / PRODUCT_SEARCH_K
            f.write(f"\n  Query: \"{r['query']}\"\n")
            f.write(f"  Keywords: {r['keywords']}\n")
            f.write(f"  Precision@{PRODUCT_SEARCH_K}: {precision:.2f} ({rel_count}/{PRODUCT_SEARCH_K})\n")
            for i, x in enumerate(top_k):
                mark = "V" if x["relevant"] else "X"
                f.write(f"    {i+1}. [{mark}] {x['title']} (dist: {x['score']:.3f})\n")

    print(f"\nResults saved to: {results_path}")


if __name__ == "__main__":
    main()