from fastapi import FastAPI, Query
from pydantic import BaseModel
import faiss
import pickle
import numpy as np
from sentence_transformers import SentenceTransformer

app = FastAPI()

# Load Sentence Transformer Model
model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

# Load Data
with open("thirukkural_data.pkl", "rb") as f:
    kural_data = pickle.load(f)

# Load FAISS indexes
english_index = faiss.read_index("thirukkural_english_index.faiss")
tamil_index = faiss.read_index("thirukkural_tamil_index.faiss")

class QueryRequest(BaseModel):
    query: str
    lang: str  # "ta" or "en"
    top_k: int = 3

@app.post("/search/")
def search_kural(req: QueryRequest):
    emb = model.encode([req.query])
    index = tamil_index if req.lang == "ta" else english_index
    D, I = index.search(np.array(emb).astype("float32"), req.top_k)

    results = [kural_data[i] for i in I[0]]
    return {"results": results}