import gradio as gr
from sentence_transformers import SentenceTransformer
import torch
import json
import torch.nn.functional as F

# --- CONFIGURATE ---
# We point to your repo to ensure it's the exact same model version
MODEL_PATH = "Ccre/gemma-embedding-model"
DATABASE_FILE = "vector_database.json"
TOP_K = 5 

# --- 1. Load Model and Database ONCE on startup ---
print("Loading model and database...")
model = SentenceTransformer(MODEL_PATH)

with open(DATABASE_FILE, 'r') as f:
    database = json.load(f)

stored_texts = [entry['text'] for entry in database]
stored_embeddings = torch.tensor([entry['embedding'] for entry in database])
print(f"Database with {len(stored_texts)} entries loaded. Ready for queries.")

# --- 2. The Search Function ---
def search_log(query):
    if not query:
        return ""
    
    query_embedding = model.encode(query, convert_to_tensor=True)
    
    cos_scores = F.cosine_similarity(query_embedding, stored_embeddings, dim=1)
    top_results = torch.topk(cos_scores, k=min(TOP_K, len(stored_texts)))
    
    results_list = []
    for score, idx in zip(top_results.values, top_results.indices):
        results_list.append(f"(Score: {score:.4f}) {stored_texts[idx]}")
    
    return "\n".join(results_list)

# --- 3. Create the Gradio Web Interface and API ---
iface = gr.Interface(
    fn=search_log,
    inputs=gr.Textbox(lines=2, label="Search Query"),
    outputs=gr.Textbox(label="Top Results"),
    title="RAG Librarian",
    description="A semantic search engine for the Nova & Zanno collaboration log."
)

iface.launch(ssr_mode=False)