import os
import numpy as np
import faiss
import gradio as gr
from pypdf import PdfReader
from sentence_transformers import SentenceTransformer
from groq import Groq

# ---------- Setup ----------
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
embedder = SentenceTransformer("all-MiniLM-L6-v2")

state = {"index": None, "chunks": None}

# ---------- Helper functions ----------
def extract_text(pdf_path):
    reader = PdfReader(pdf_path)
    text = ""
    for page in reader.pages:
        page_text = page.extract_text()
        if page_text:
            text += page_text + "\n"
    return text

def chunk_text(text, chunk_size=300, overlap=50):
    words = text.split()
    chunks = []
    i = 0
    while i < len(words):
        chunk = " ".join(words[i:i + chunk_size])
        chunks.append(chunk)
        i += chunk_size - overlap
    return chunks

def build_index(chunks):
    embeddings = embedder.encode(chunks, show_progress_bar=False)
    embeddings = np.array(embeddings).astype("float32")
    index = faiss.IndexFlatL2(embeddings.shape[1])
    index.add(embeddings)
    return index

def retrieve(question, index, chunks, k=4):
    q_emb = embedder.encode([question]).astype("float32")
    distances, indices = index.search(q_emb, k)
    return [chunks[i] for i in indices[0] if i < len(chunks)]

def generate_answer(question, context_chunks):
    context = "\n\n".join(context_chunks)
    prompt = f"""Answer the question based only on the context below.
If the answer isn't in the context, say you don't know.

Context:
{context}

Question: {question}

Answer:"""

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

# ---------- App functions ----------
def process_pdf(pdf_file):
    if pdf_file is None:
        return "Please upload a PDF first."
    text = extract_text(pdf_file.name)
    if not text.strip():
        return "Couldn't read any text from this PDF (it may be scanned images)."
    chunks = chunk_text(text)
    index = build_index(chunks)
    state["index"] = index
    state["chunks"] = chunks
    return f"✅ PDF processed into {len(chunks)} chunks. You can ask questions now."

def answer_question(question, history):
    if state["index"] is None:
        return history + [[question, "Please upload and process a PDF first."]]
    if not question or not question.strip():
        return history
    relevant_chunks = retrieve(question, state["index"], state["chunks"])
    answer = generate_answer(question, relevant_chunks)
    return history + [[question, answer]]

# ---------- UI ----------
with gr.Blocks(title="Chat with your PDF") as demo:
    gr.Markdown("# 📄 Chat with your PDF\nUpload a PDF, then ask questions about it.")

    with gr.Row():
        pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
        upload_btn = gr.Button("Process PDF", variant="primary")
    status = gr.Textbox(label="Status", interactive=False)

    chatbot = gr.Chatbot(label="Conversation")
    question = gr.Textbox(label="Ask a question", placeholder="Type your question and press Enter")

    upload_btn.click(process_pdf, inputs=pdf_input, outputs=status)
    question.submit(answer_question, inputs=[question, chatbot], outputs=chatbot)
    question.submit(lambda: "", None, question)

demo.launch()