import os
import gradio as gr
import tempfile

from groq import Groq

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings

# -----------------------------
# Groq Client
# -----------------------------
client = Groq(api_key=os.environ.get("RAG_API_KEY"))

# -----------------------------
# Load & Process Document
# -----------------------------
def load_document(file_path):
    if file_path.endswith(".pdf"):
        loader = PyPDFLoader(file_path)
    else:
        loader = TextLoader(file_path, encoding="utf-8")
    return loader.load()

def build_vector_db(docs):
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=100
    )
    chunks = splitter.split_documents(docs)

    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )

    db = FAISS.from_documents(chunks, embeddings)
    return db

# -----------------------------
# RAG Question Answering
# -----------------------------
def answer_question(file, question):
    if file is None or question.strip() == "":
        return "❌ Please upload some document and enter a question."

    with tempfile.NamedTemporaryFile(delete=False) as tmp:
        tmp.write(file.read())
        file_path = tmp.name

    documents = load_document(file_path)
    vector_db = build_vector_db(documents)

    relevant_docs = vector_db.similarity_search(question, k=4)
    context = "\n\n".join([doc.page_content for doc in relevant_docs])

    # ✅ UPDATED PROMPT (Document + General Knowledge)
    prompt = f"""
You are a knowledgeable and helpful assistant.

Use the provided document context as your PRIMARY source.
If the document does not fully answer the question, you MAY use your general knowledge to give additional relevant information.

Rules:
- Prefer the document context whenever possible
- If you use information outside the document, clearly mention that it is general knowledge
- Do not invent facts

Document Context:
{context}

User Question:
{question}

Answer:
"""

    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[{"role": "user", "content": prompt}],
    )

    return response.choices[0].message.content

# -----------------------------
# Gradio UI
# -----------------------------
with gr.Blocks() as demo:
    gr.Markdown("# 📄 RAG Document Question Answering App")
    gr.Markdown(
        "Upload a document (PDF or TXT) and ask questions about it.\n"
        "The app prioritizes document content but can also use general knowledge when needed."
    )

    with gr.Row():
        file_input = gr.File(label="Upload Document (PDF / TXT)")
        question_input = gr.Textbox(
            label="Ask a Question",
            placeholder="What is this document about?"
        )

    answer_output = gr.Textbox(
        label="Answer",
        lines=10
    )

    ask_btn = gr.Button("Ask Question")

    ask_btn.click(
        fn=answer_question,
        inputs=[file_input, question_input],
        outputs=answer_output
    )

demo.launch()