import os
import fitz  # PyMuPDF
import gradio as gr
import requests
from dotenv import load_dotenv
from sentence_transformers import SentenceTransformer, util

# Load environment variables
load_dotenv()
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"
MODEL_NAME = "llama3-8b-8192"

# Load sentence-transformer model
embed_model = SentenceTransformer('all-MiniLM-L6-v2')

SYSTEM_PROMPT = """
You are a helpful AI assistant that answers questions using provided PDF documents.
Always base your answers on the context and include page numbers when possible.
"""

# Extract text from PDF with page numbers
def extract_text_with_pages(file):
    doc = fitz.open(file.name)
    chunks = []
    for page_num in range(len(doc)):
        page = doc.load_page(page_num)
        text = page.get_text()
        if text.strip():
            chunks.append({"text": text.strip(), "page": page_num + 1})
    return chunks

# Chunking (simple, per page)
def prepare_embeddings(chunks):
    texts = [chunk["text"] for chunk in chunks]
    embeddings = embed_model.encode(texts, convert_to_tensor=True)
    return embeddings

# Retrieve top chunks
def retrieve_top_chunks(query, chunks, embeddings, top_k=3):
    query_embedding = embed_model.encode(query, convert_to_tensor=True)
    hits = util.semantic_search(query_embedding, embeddings, top_k=top_k)[0]
    results = []
    for hit in hits:
        matched_chunk = chunks[hit["corpus_id"]]
        results.append(f"[Page {matched_chunk['page']}]: {matched_chunk['text']}")
    return results

# Query GROQ LLM
def query_groq(question, context, chat_history):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for user, bot in chat_history:
        messages.append({"role": "user", "content": user})
        messages.append({"role": "assistant", "content": bot})
    messages.append({"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"})

    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json"
    }

    response = requests.post(GROQ_API_URL, headers=headers, json={
        "model": MODEL_NAME,
        "messages": messages,
        "temperature": 0.7
    })

    if response.status_code == 200:
        return response.json()["choices"][0]["message"]["content"]
    else:
        return f"Error {response.status_code}: {response.text}"

# Full RAG pipeline
def rag_chatbot(files, question, history):
    all_chunks = []
    for file in files:
        all_chunks.extend(extract_text_with_pages(file))
    
    embeddings = prepare_embeddings(all_chunks)
    top_chunks = retrieve_top_chunks(question, all_chunks, embeddings)
    context = "\n\n".join(top_chunks)
    response = query_groq(question, context, history)
    history.append((question, response))
    return "", history

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## 📄 RAG PDF Chatbot (Groq + SentenceTransformers)")
    gr.Markdown("Upload PDFs and ask questions. The assistant will use document context and show page numbers in answers.")

    pdf_files = gr.File(file_types=[".pdf"], file_count="multiple", label="Upload your PDF files")
    question = gr.Textbox(label="Ask a question about the documents")
    chatbot = gr.Chatbot()
    clear_btn = gr.Button("Clear Chat")
    state = gr.State([])

    question.submit(rag_chatbot, [pdf_files, question, state], [question, chatbot])
    clear_btn.click(lambda: ([], []), None, [chatbot, state])

demo.launch()