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

# ==============================
# CONFIG
# ==============================

EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
GROQ_MODEL = "llama-3.1-8b-instant"

# Load embedding model
embedder = SentenceTransformer(EMBEDDING_MODEL)

# Load Groq client
client = Groq(
    api_key=os.environ.get("GROQ_API_KEY")
)

# Global storage
vector_store = None
stored_chunks = []


# ==============================
# PDF PROCESSING
# ==============================
pdf_input = gr.File(file_types=[".pdf"], type="filepath")

def extract_text_from_pdf(pdf_file):
    reader = PdfReader(pdf_file)
    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=500, overlap=100):
    chunks = []
    start = 0

    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start = end - overlap

    return chunks


# ==============================
# CREATE VECTOR STORE
# ==============================

def create_vector_store(chunks):
    global vector_store

    embeddings = embedder.encode(chunks)
    embeddings = np.array(embeddings).astype("float32")  # IMPORTANT

    dimension = embeddings.shape[1]

    index = faiss.IndexFlatL2(dimension)
    index.add(embeddings)

    vector_store = index


# ==============================
# RETRIEVAL
# ==============================

def retrieve_chunks(query, k=3):
    query_embedding = embedder.encode([query])
    query_embedding = np.array(query_embedding).astype("float32")

    distances, indices = vector_store.search(query_embedding, k)

    results = []
    for i in indices[0]:
        if i < len(stored_chunks):
            results.append(stored_chunks[i])

    return "\n\n".join(results)


# ==============================
# GROQ RESPONSE
# ==============================

def generate_answer(context, question):
    try:
        prompt = f"""
You are a helpful AI assistant.
Answer ONLY from the provided context.
If the answer is not in the context, say "Not found in document."

Context:
{context}

Question:
{question}
"""

        chat_completion = client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}],
            model=GROQ_MODEL,
        )

        return chat_completion.choices[0].message.content

    except Exception as e:
        return f"Error generating answer: {str(e)}"


# ==============================
# MAIN PIPELINE
# ==============================

def process_pdf(pdf):
    global stored_chunks

    if pdf is None:
        return "Please upload a PDF."

    text = extract_text_from_pdf(pdf)

    if len(text.strip()) == 0:
        return "No readable text found in PDF."

    stored_chunks = chunk_text(text)

    create_vector_store(stored_chunks)

    return "PDF processed successfully! You can now ask questions."


def answer_question(question):
    if vector_store is None:
        return "Please upload and process a PDF first."

    if not question.strip():
        return "Please enter a valid question."

    context = retrieve_chunks(question)

    if not context:
        return "No relevant information found in document."

    answer = generate_answer(context, question)
    return answer


# ==============================
# GRADIO UI
# ==============================

with gr.Blocks(
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        font-family: 'Inter', sans-serif;
    }

    /* Rounded buttons */
    button {
        border-radius: 999px !important;
        padding: 10px 18px !important;
        font-weight: 600;
    }

    /* Card styling */
    .card {
        background: #ffffff;
        padding: 20px;
        border-radius: 16px;
        box-shadow: 0px 4px 20px rgba(0,0,0,0.05);
    }

    /* Input fields */
    textarea, input {
        border-radius: 12px !important;
    }
    """
) as demo:

    gr.Markdown(
        """
        <h1 style='text-align: center;'>📄 RAG PDF Assistant</h1>
        <p style='text-align: center; color: gray;'>
        Upload your PDF and ask intelligent questions
        </p>
        """
    )

    with gr.Row():
        # LEFT CARD (UPLOAD)
        with gr.Column(scale=1):
            with gr.Group(elem_classes="card"):
                pdf_input = gr.File(label="📂 Upload PDF", file_types=[".pdf"])
                upload_button = gr.Button("⚙️ Process PDF", variant="primary")
                status_output = gr.Textbox(label="Status", lines=2)

        # RIGHT CARD (QA)
        with gr.Column(scale=2):
            with gr.Group(elem_classes="card"):
                question_input = gr.Textbox(
                    label="💬 Ask a question",
                    placeholder="Type your question here..."
                )
                ask_button = gr.Button("🚀 Get Answer", variant="primary")
                answer_output = gr.Textbox(label="Answer", lines=8)

    gr.Markdown(
        "<p style='text-align:center; font-size:12px; color:gray;'>Built with ❤️ using RAG + LLM</p>"
    )

    upload_button.click(process_pdf, inputs=pdf_input, outputs=status_output)
    ask_button.click(answer_question, inputs=question_input, outputs=answer_output)

demo.launch()