import gradio as gr
from sentence_transformers import SentenceTransformer, CrossEncoder
import chromadb
from rank_bm25 import BM25Okapi
from groq import Groq
import json, os

# Load corpus
with open("corpus.json", "r", encoding="utf-8") as f:
    final_corpus = json.load(f)

# Setup models
embed_model = SentenceTransformer("intfloat/multilingual-e5-base")
chroma_client = chromadb.Client()
collection = chroma_client.create_collection(name="hindi_finrag")

for doc in final_corpus:
    embedding = embed_model.encode(doc["text"]).tolist()
    collection.add(ids=[doc["id"]], embeddings=[embedding], documents=[doc["text"]])

tokenized_corpus = [doc["text"].split() for doc in final_corpus]
bm25 = BM25Okapi(tokenized_corpus)
reranker = CrossEncoder("cross-encoder/mmarco-mMiniLMv2-L12-H384-v1")

GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
groq_client = Groq(api_key=GROQ_API_KEY)

def rrf_fusion(dense_ids, sparse_ids, k=60):
    scores = {}
    for rank, doc_id in enumerate(dense_ids):
        scores[doc_id] = scores.get(doc_id, 0) + 1/(k+rank+1)
    for rank, doc_id in enumerate(sparse_ids):
        scores[doc_id] = scores.get(doc_id, 0) + 1/(k+rank+1)
    return sorted(scores.keys(), key=lambda x: scores[x], reverse=True)

def hybrid_retrieve_reranked(question, top_k=3):
    query_emb = embed_model.encode(question).tolist()
    dense = collection.query(query_embeddings=[query_emb], n_results=len(final_corpus))
    dense_ids = dense["ids"][0]
    bm25_scores = bm25.get_scores(question.split())
    sparse_ids = [final_corpus[i]["id"] for i in sorted(range(len(bm25_scores)), key=lambda i: bm25_scores[i], reverse=True)]
    fused = rrf_fusion(dense_ids, sparse_ids)
    id_to_text = {doc["id"]: doc["text"] for doc in final_corpus}
    candidates = [id_to_text[i] for i in fused[:10]]
    scores = reranker.predict([(question, doc) for doc in candidates])
    ranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
    return [doc for doc, _ in ranked[:top_k]]

def rag_query(question):
    if not question.strip():
        return "कृपया एक प्रश्न लिखें।", ""
    contexts = hybrid_retrieve_reranked(question)
    context_text = " ".join(contexts)[:1200]
    prompt = f"""आप एक वित्तीय सहायक हैं। नीचे दिए गए संदर्भ के आधार पर प्रश्न का उत्तर दें।

संदर्भ:
{context_text}

प्रश्न: {question}

उत्तर (हिंदी में, संक्षिप्त):"""
    response = groq_client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.1
    )
    answer = response.choices[0].message.content
    sources = "\n\n".join([f"स्रोत {i+1}: {c[:200]}..." for i, c in enumerate(contexts)])
    return answer, sources

demo = gr.Interface(
    fn=rag_query,
    inputs=gr.Textbox(label="अपना प्रश्न हिंदी में लिखें", placeholder="उदाहरण: पीएम-किसान योजना क्या है?"),
    outputs=[gr.Textbox(label="उत्तर"), gr.Textbox(label="स्रोत दस्तावेज़")],
    title="HindiFinRAG — Hindi Financial RAG System",
    description="भारतीय वित्तीय और सरकारी योजनाओं के बारे में हिंदी में प्रश्न पूछें।"
)

demo.launch()