import os
import gradio as gr
import faiss
import numpy as np
import PyPDF2
from sentence_transformers import SentenceTransformer
import requests

# 🔐 Use environment variable for Groq API key
GROQ_API_KEY = os.getenv("GROQ_API_KEY", "")
GROQ_MODEL = "llama3-8b-8192"

# 🧠 Embedding Model
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

# 🗃️ Vector DB
dimension = 384
index = faiss.IndexFlatL2(dimension)
text_chunks = []

def extract_text_from_pdf(pdf_file):
    reader = PyPDF2.PdfReader(pdf_file)
    text = ""
    for page in reader.pages:
        text += page.extract_text() or ""
    return text

def chunk_text(text, chunk_size=500):
    sentences = text.split('. ')
    chunks, chunk = [], ""
    for sentence in sentences:
        if len(chunk) + len(sentence) < chunk_size:
            chunk += sentence + ". "
        else:
            chunks.append(chunk.strip())
            chunk = sentence + ". "
    if chunk:
        chunks.append(chunk.strip())
    return chunks

def embed_and_store(chunks):
    global text_chunks
    text_chunks = chunks
    embeddings = embedding_model.encode(chunks)
    index.add(np.array(embeddings))

def retrieve_context(query, top_k=3):
    query_vector = embedding_model.encode([query])
    distances, indices = index.search(np.array(query_vector), top_k)
    return "\n".join([text_chunks[i] for i in indices[0]])

def format_prompt(context, question):
    system_msg = "You are a helpful research assistant who answers questions using only the uploaded document."
    user_msg = f"Document Context:\n{context}\n\nQuestion: {question}\nAnswer:"
    return [{"role": "system", "content": system_msg},
            {"role": "user", "content": user_msg}]

def call_groq_api(messages):
    url = "https://api.groq.com/openai/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json"
    }
    data = {
        "model": GROQ_MODEL,
        "messages": messages,
        "temperature": 0.3
    }
    response = requests.post(url, headers=headers, json=data)
    return response.json()['choices'][0]['message']['content']

def upload_file(pdf):
    text = extract_text_from_pdf(pdf)
    chunks = chunk_text(text)
    embed_and_store(chunks)
    return "✅ Document processed. You may now ask questions."

def answer_question(question):
    if not text_chunks:
        return "❌ Please upload and process a document first."
    context = retrieve_context(question)
    messages = format_prompt(context, question)
    return call_groq_api(messages)

with gr.Blocks() as rag_ui:
    gr.Markdown("## 📄 RAG Assistant with LLaMA3 (Groq)")
    
    with gr.Row():
        pdf_input = gr.File(label="Upload PDF")
        upload_button = gr.Button("Process Document")
    
    status_output = gr.Textbox(label="Status")
    upload_button.click(upload_file, inputs=pdf_input, outputs=status_output)

    gr.Markdown("### ❓ Ask a Question from the Uploaded PDF")
    question_input = gr.Textbox(label="Your Question")
    answer_output = gr.Textbox(label="Answer", lines=5)

    ask_button = gr.Button("Get Answer")
    ask_button.click(answer_question, inputs=question_input, outputs=answer_output)

rag_ui.launch()