# relationship_retriever.py
import os, pickle, logging
import faiss
from sentence_transformers import SentenceTransformer

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

RELATIONS = "relations"
REL_INDEX = f"{RELATIONS}/got_rels.faiss"
REL_DATA = f"{RELATIONS}/got_rels_meta.pkl"

logger.info("Loading relationship FAISS index...")
rel_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2",
                                device="cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu")
rel_index = faiss.read_index(REL_INDEX)
with open(REL_DATA, "rb") as f:
    rel_data = pickle.load(f)
name_map = rel_data["name_map"]

def batch_relationships(questions, top_k=3):
    batch_results = []
    for q in questions:
        q_upper = q.upper()
        candidates = []
        for variant in name_map.keys():
            if len(variant) < 3: continue
            if variant in q_upper or variant.replace(" ","") in q_upper.replace(" ",""):
                candidates.append(name_map[variant])
        candidates = list(dict.fromkeys(candidates))[:2]
        if not candidates:
            batch_results.append(["No known character relationships found"])
            continue

        query = f"Relationships of {' and '.join(candidates)} in Game of Thrones books"
        q_vec = rel_model.encode([query], normalize_embeddings=True, show_progress_bar=False).astype("float32")
        D, I = rel_index.search(q_vec, top_k*2)
        results = []
        seen = set()
        for idx in I[0]:
            if idx == -1: continue
            sent = rel_data["sentences"][idx]
            char = rel_data["metadata"][idx]["display_name"]
            if char not in seen:
                results.append(sent)
                seen.add(char)
            if len(results) >= top_k: break
        batch_results.append(results if results else ["No confirmed relationships found"])
    return batch_results