print("DEBUG: Starting Retrieve.py")

import faiss
import numpy as np
import pickle
import os
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_core.documents import Document
import streamlit as st
from config import VECTORSTORE_PATHS, SIMILARITY_THRESHOLD, MAX_CHUNKS

print("DEBUG: Retrieve.py imports completed")

def normalize(v):
    v = np.array(v)
    norm = np.linalg.norm(v)
    return v if norm == 0 else (v / norm)

def retrieve_chunks_from_multiple_vdbs(query, selected_dbs, SIMILARITY_THRESHOLD, MAX_CHUNKS):
    hf_token = os.getenv("HF_TOKEN")
    embedding_model = HuggingFaceEmbeddings(
        model_name="intfloat/e5-large-v2",
        model_kwargs={"token": hf_token} if hf_token else {}
    )
    formatted_query = f"query: {query.strip()}"
    query_vector = embedding_model.embed_query(formatted_query)
    query_vector = normalize(query_vector).astype("float32").reshape(1, -1)

    combined_results = []
    for db_key in selected_dbs:
        vectorstore_path = VECTORSTORE_PATHS.get(db_key)
        if not vectorstore_path:
            continue

        index_path = f"{vectorstore_path}/faiss.index"
        docstore_path = f"{vectorstore_path}/documents.pkl"

        index = faiss.read_index(index_path)
        with open(docstore_path, "rb") as f:
            documents: list[Document] = pickle.load(f)

        scores, indices = index.search(query_vector, MAX_CHUNKS * 2)
        for score, idx in zip(scores[0], indices[0]):
            if idx == -1 or idx >= len(documents):
                continue
            doc = documents[idx]
            if score >= SIMILARITY_THRESHOLD:
                combined_results.append((score, doc))

    combined_results.sort(key=lambda x: -x[0])

    selected_chunks = []
    references = []
    seen_metadata = set()

    for score, doc in combined_results:
        source = doc.metadata.get("source", "Unknown")
        page = doc.metadata.get("page", "Unknown")
        if (source, page) in seen_metadata:
            continue
        seen_metadata.add((source, page))

        reference = f"[Chunk {len(selected_chunks)+1}] Source: {source}, Page: {page}, Similarity: {score:.2f}"
        selected_chunks.append(f"{doc.page_content}\n{reference}")
        references.append(reference)

        if len(selected_chunks) >= MAX_CHUNKS:
            break

    return selected_chunks, references