from langchain_community.llms import HuggingFaceHub
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.chains import RetrievalQA
import os
import requests
import gradio as gr

# Configuration
HF_TOKEN = os.environ.get("HF_TOKEN", "")
CLOUDFLARE_API = "https://notary-662-sbz.pages.dev/api/db/chats"

# 1. Setup Llama 3 via Hugging Face Inference API
llm = HuggingFaceHub(
    repo_id="meta-llama/Meta-Llama-3-8B-Instruct",
    huggingfacehub_api_token=HF_TOKEN,
    model_kwargs={"temperature": 0.7, "max_new_tokens": 512}
)

# 2. Setup Persian-capable Embeddings
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

def query_rag_system(user_query, chat_id="default"):
    # Load the 100 docs index (assuming it's saved locally in the Space)
    try:
        vector_store = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
        qa_chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=vector_store.as_retriever())
        
        # Get response from Llama 3 + RAG
        result = qa_chain.run(user_query)
        
        # 3. Sync with Cloudflare D1
        chat_data = {
            "id": chat_id,
            "title": user_query[:30],
            "docType": "notary_rag_llama",
            "messages": [
                {"role": "user", "text": user_query},
                {"role": "model", "text": result}
            ]
        }
        requests.post(CLOUDFLARE_API, json=chat_data)
        
        return result
    except Exception as e:
        return f"خطا در اتصال به بانک اسناد: {str(e)}"

# Gradio Interface
iface = gr.Interface(
    fn=query_rag_system,
    inputs="text",
    outputs="text",
    title="Notary Llama-3 RAG Engine",
    description="این سیستم مستقیماً به Cloudflare و ۱۰۰ فایل PDF محضر متصل است."
)

if __name__ == "__main__":
    iface.launch()