import gradio as gr
import openai
import os
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.schema import Document
from sentence_transformers import SentenceTransformer
from datasets import load_dataset
import nltk

# ✅ Load the Sentence Transformer Embedding Model
model_name = "sentence-transformers/all-MiniLM-L6-v2"
embedding_model = HuggingFaceEmbeddings(model_name=model_name)

# ✅ Set OpenAI API Key
openai.api_key = os.getenv("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
")

# ✅ Download NLTK Tokenizer
nltk.download('punkt')

# ✅ Load and Chunk Dataset
def chunk_documents(documents, max_chunk_size=500):
    chunks = []
    for doc in documents:
        sentences = nltk.sent_tokenize(doc)
        current_chunk = ""
        for sentence in sentences:
            if len(current_chunk) + len(sentence) <= max_chunk_size:
                current_chunk += sentence + " "
            else:
                chunks.append(current_chunk.strip())
                current_chunk = sentence + " "
        if current_chunk:
            chunks.append(current_chunk.strip())
    return chunks

# ✅ Load Dataset and Prepare ChromaDB
dataset = load_dataset("rungalileo/ragbench", "techqa")  # Example dataset
original_documents = dataset['train']['documents']
chunked_documents = chunk_documents(original_documents)

persist_directory = "chroma_db_directory"
documents = [Document(page_content=chunk) for chunk in chunked_documents]

# ✅ Initialize ChromaDB
vectordb = Chroma.from_documents(
    documents=documents,
    embedding=embedding_model,
    persist_directory=persist_directory
)
vectordb.persist()

# ✅ Function to Retrieve Relevant Documents
def retrieve_documents(question, k=5):
    docs = vectordb.similarity_search(question, k=k)
    if not docs:
        return ["⚠️ No relevant documents found. Try a different query."]
    return [doc.page_content for doc in docs]

# ✅ Function to Generate AI Response
def generate_response(question, context):
    if not context or "No relevant documents found." in context:
        return "No relevant context available. Try a different query."

    full_prompt = f"Context: {context}\n\nQuestion: {question}"

    try:
        response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=[
                {"role": "system", "content": "You are an AI assistant that answers user queries based on the given context."},
                {"role": "user", "content": full_prompt}
            ],
            max_tokens=300,
            temperature=0.7
        )
        return response['choices'][0]['message']['content'].strip()
    except Exception as e:
        return f"Error generating response: {str(e)}"

# ✅ Full RAG Pipeline
def rag_pipeline(question):
    retrieved_docs = retrieve_documents(question, k=5)
    context = " ".join(retrieved_docs)
    response = generate_response(question, context)
    return response, "\n\n".join(retrieved_docs)

# ✅ Gradio UI Interface
iface = gr.Interface(
    fn=rag_pipeline,
    inputs=gr.Textbox(label="Enter your question"),
    outputs=[
        gr.Textbox(label="Generated Response"),
        gr.Textbox(label="Retrieved Documents")
    ],
    title="RAG-Based Question Answering System",
    description="Enter a question and retrieve relevant documents with AI-generated response."
)

# ✅ Launch the Gradio App
if __name__ == "__main__":
    iface.launch()