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

# -----------------------------
# Initialize Models
# -----------------------------
embedder = SentenceTransformer("all-MiniLM-L6-v2")

GROQ_API_KEY = os.getenv("Rag")
client = Groq(api_key=GROQ_API_KEY) if GROQ_API_KEY else None

# -----------------------------
# Global Storage
# -----------------------------
documents = []
embeddings = None

# -----------------------------
# PDF Processing
# -----------------------------
def read_pdf(file):
    try:
        reader = PdfReader(file.name)
        text = ""
        for page in reader.pages:
            content = page.extract_text()
            if content:
                text += content
        return text
    except Exception as e:
        return f"Error reading PDF: {str(e)}"


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 += chunk_size - overlap

    return chunks


# -----------------------------
# Create Embeddings
# -----------------------------
def create_embeddings(chunks):
    global documents, embeddings

    documents = chunks
    embeddings = embedder.encode(chunks)
    embeddings = np.array(embeddings)


# -----------------------------
# Cosine Similarity Retrieval
# -----------------------------
def cosine_similarity(a, b):
    return np.dot(a, b.T) / (np.linalg.norm(a) * np.linalg.norm(b, axis=1))


def retrieve(query, k=3, threshold=0.3):
    global embeddings

    if embeddings is None:
        return [], None

    query_embedding = embedder.encode([query])[0]

    sims = cosine_similarity(query_embedding, embeddings)

    top_k_idx = np.argsort(sims)[-k:][::-1]

    relevant_chunks = []
    scores = []

    for i in top_k_idx:
        if sims[i] > threshold:
            relevant_chunks.append(documents[i])
            scores.append(sims[i])

    # Confidence
    confidence = None
    if scores:
        avg = np.mean(scores)
        if avg > 0.7:
            confidence = "High"
        elif avg > 0.5:
            confidence = "Medium"
        else:
            confidence = "Low"

    return relevant_chunks, confidence


# -----------------------------
# Groq LLM
# -----------------------------
def ask_groq(context_chunks, question):
    if client is None:
        return "Error: Please set GROQ_API_KEY in Hugging Face Secrets."

    context = "\n".join(context_chunks)

    prompt = f"""
You are an intelligent assistant.

Rules:
1. If answer is clearly in context → answer normally.
2. If related but not exact → say:
   "This is not explicitly mentioned in the document, but based on related context..."
3. If irrelevant → say:
   "The document does not contain information related to this question."

Context:
{context}

Question:
{question}
"""

    try:
        response = client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}],
            model="llama-3.3-70b-versatile",
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Groq API Error: {str(e)}"


# -----------------------------
# Main Functions
# -----------------------------
def process_pdf(file):
    if file is None:
        return "Please upload a PDF."

    text = read_pdf(file)

    if not text or "Error" in text:
        return text

    chunks = chunk_text(text)
    create_embeddings(chunks)

    return f"✅ PDF processed successfully! Chunks: {len(chunks)}"


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

    context_chunks, confidence = retrieve(question)

    if not context_chunks:
        return "The document does not contain information related to this question."

    answer = ask_groq(context_chunks, question)

    if confidence:
        answer = f"(Confidence: {confidence})\n\n{answer}"

    return answer


# -----------------------------
# Gradio UI
# -----------------------------
with gr.Blocks() as demo:
    gr.Markdown("## 📄 RAG PDF Q&A (Groq + HuggingFace Ready)")

    file_input = gr.File(label="Upload PDF")
    upload_btn = gr.Button("Process PDF")
    status = gr.Textbox(label="Status")

    question = gr.Textbox(label="Ask a question")
    answer = gr.Textbox(label="Answer")

    upload_btn.click(process_pdf, inputs=file_input, outputs=status)
    question.submit(answer_question, inputs=question, outputs=answer)


# -----------------------------
# Launch
# -----------------------------
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)