import os
import gradio as gr
import fitz
import chromadb
from sentence_transformers import SentenceTransformer
from groq import Groq

# -----------------------------
# API KEY
# -----------------------------

GROQ_API_KEY = os.getenv("GROQ_API_KEY")

client = Groq(api_key=GROQ_API_KEY)

# -----------------------------
# EMBEDDING MODEL
# -----------------------------

embedding_model = None

def get_embedding_model():

    global embedding_model

    if embedding_model is None:

        embedding_model = SentenceTransformer(
            "sentence-transformers/all-MiniLM-L6-v2"
        )

    return embedding_model

# -----------------------------
# CHROMADB
# -----------------------------



chroma_client = chromadb.PersistentClient(path="./chroma_db")

try:
    chroma_client.delete_collection("pdf_rag_collection")
except:
    pass

collection = chroma_client.create_collection(
    name="pdf_rag_collection"
)

# -----------------------------
# CHUNKING
# -----------------------------

def chunk_text(text, chunk_size=800, overlap=200):

    chunks = []

    start = 0

    while start < len(text):

        end = start + chunk_size

        chunks.append(text[start:end])

        start += chunk_size - overlap

    return chunks

# -----------------------------
# PROCESS PDF
# -----------------------------

def process_pdf(pdf_file):

    text = ""

    doc = fitz.open(pdf_file.name)

    for page in doc:
        text += page.get_text()

    chunks = chunk_text(text)
    print("TOTAL CHUNKS:", len(chunks))
    print("FIRST CHUNK:", chunks[0][:300])

    try:
        collection.delete(where={})
    except:
        pass

    for i, chunk in enumerate(chunks):

        embedding = get_embedding_model().encode(chunk).tolist()

        collection.add(
            ids=[str(i)],
            embeddings=[embedding],
            documents=[chunk]
        )

    return "PDF processed successfully!"

# -----------------------------
# RETRIEVAL
# -----------------------------

def retrieve_context(question):

    query_embedding = get_embedding_model().encode(question).tolist()

    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=3
    )

    docs = results["documents"][0]
    print(results)

    return "\n".join(docs)

# -----------------------------
# CHATBOT
# -----------------------------

def chatbot(question):

    context = retrieve_context(question)

    prompt = f"""
Answer ONLY using the context below.

If answer is not present,
say:
'I could not find this information in the document.'

Context:
{context}

Question:
{question}
"""

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

    return response.choices[0].message.content

# -----------------------------
# UI
# -----------------------------

with gr.Blocks() as demo:

    gr.Markdown("# 📄 PDF RAG Chatbot")

    pdf_input = gr.File(label="Upload PDF")

    process_btn = gr.Button("Process PDF")

    status = gr.Textbox(label="Status")

    process_btn.click(
        process_pdf,
        inputs=pdf_input,
        outputs=status
    )

    question = gr.Textbox(label="Ask Question")

    ask_btn = gr.Button("Ask")

    answer = gr.Textbox(label="Answer")

    ask_btn.click(
        chatbot,
        inputs=question,
        outputs=answer
    )
demo.launch(server_name="0.0.0.0", server_port=7860)