import os
import gradio as gr
import requests
from bs4 import BeautifulSoup
from PyPDF2 import PdfReader
from docx import Document
from pptx import Presentation
import numpy as np
import faiss
from groq import Groq
import re

# ---------------- Groq client ----------------
client = Groq(api_key=os.environ.get("MY_API"))

# ---------------- FAISS setup ----------------
dimension = 768
index = faiss.IndexFlatL2(dimension)
texts_storage = []

# ---------------- Chunking ----------------
def chunk_text(text, chunk_size=500):
    words = text.split()
    chunks = []
    for i in range(0, len(words), chunk_size):
        chunks.append(" ".join(words[i:i+chunk_size]))
    return chunks

# ---------------- Fake Embedding (replace with real later) ----------------
def embed(text):
    return np.random.rand(dimension).astype("float32")

# ---------------- Retrieval ----------------
def retrieve_context(query, top_k=5):
    q_emb = embed(query)
    distances, positions = index.search(np.array([q_emb]), top_k)
    results = [texts_storage[p] for p in positions[0] if p < len(texts_storage)]
    return "\n".join(results)

# ---------------- Document Extractors ----------------
def extract_pdf(file):
    reader = PdfReader(file)
    text = ""
    for page in reader.pages:
        text += page.extract_text() or ""
    return text

def extract_docx(file):
    doc = Document(file)
    return "\n".join([p.text for p in doc.paragraphs])

def extract_pptx(file):
    prs = Presentation(file)
    text = ""
    for slide in prs.slides:
        for shape in slide.shapes:
            if hasattr(shape, "text"):
                text += shape.text + "\n"
    return text

def extract_text_from_url(url):
    try:
        r = requests.get(url, timeout=10)
        r.raise_for_status()
        soup = BeautifulSoup(r.text, "html.parser")
        for s in soup(["script", "style"]):
            s.decompose()
        return soup.get_text(separator="\n")
    except Exception as e:
        return f"Error fetching URL: {e}"

# ---------------- Build FAISS Index with Chunks ----------------
def build_index(text):
    chunks = chunk_text(text)
    for ch in chunks:
        emb = embed(ch)
        index.add(np.array([emb]))
        texts_storage.append(ch)
    return "Index built successfully"

# ---------------- Generate Summary, MCQs, Short QA separately ----------------
def generate_summary(text):
    prompt = f"Write a detailed summary of the following content:\n\n{text}"
    response = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama-3.1-8b-instant"
    )
    return response.choices[0].message.content

def generate_mcqs(text, num=20):
    prompt = f"Create {num} multiple choice questions with answers based on the text below:\n\n{text}"
    response = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama-3.1-8b-instant"
    )
    return response.choices[0].message.content

def generate_shortqa(text, num=20):
    prompt = f"Create {num} short questions and answers based on the text below:\n\n{text}"
    response = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama-3.1-8b-instant"
    )
    return response.choices[0].message.content

# ---------------- Ask Question with RAG ----------------
def ask_question(question):
    context = retrieve_context(question, top_k=5)

    prompt = f"""
Answer based on the context below.
If context is irrelevant, answer general knowledge.
CONTEXT:
{context}
Question: {question}
"""
    response = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama-3.1-8b-instant"
    )
    return response.choices[0].message.content

# ---------------- Split Output into Sections ----------------
def split_analysis(text):
    summary = mcqs = short_qa = ""

    summary_match = re.search(r"\*\*Summary\*\*(.*?)(?=\*\*15 Multiple Choice Questions)", text, re.DOTALL)
    mcqs_match = re.search(r"\*\*15 Multiple Choice Questions\*\*(.*?)(?=\*\*15 Short Q/A)", text, re.DOTALL)
    short_match = re.search(r"\*\*15 Short Q/A\*\*(.*)", text, re.DOTALL)

    if summary_match: summary = summary_match.group(1).strip()
    if mcqs_match: mcqs = mcqs_match.group(1).strip()
    if short_match: short_qa = short_match.group(1).strip()

    return summary, mcqs, short_qa
    # -------- Button Functions --------
def process_documents(files):
    # RESET FAISS + STORAGE
    index.reset()
    texts_storage.clear()

    text = ""
    for f in files:
        n = f.name.lower()
        if n.endswith(".pdf"):
            text += extract_pdf(f)
        elif n.endswith(".docx"):
            text += extract_docx(f)
        elif n.endswith(".pptx"):
            text += extract_pptx(f)
        else:
            text += f.read().decode("utf-8")

    build_index(text)

    summary = generate_summary(text)
    mcqs = generate_mcqs(text, num=20)
    shortqa = generate_shortqa(text, num=20)

    return summary, mcqs, shortqa


def process_website(url):
    index.reset()
    texts_storage.clear()

    text = extract_text_from_url(url)
    build_index(text)

    summary = generate_summary(text)
    mcqs = generate_mcqs(text, num=20)
    shortqa = generate_shortqa(text, num=20)

    return summary, mcqs, shortqa


custom_css = """
/* ---------- Global Dark Theme ---------- */
body, .gradio-container {
    font-family: 'Inter', sans-serif;
    background: #0d0d0d !important;
    color: #d6e2f0 !important;
}
/* ---------- Page Title ---------- */
h1, h2, h3, h4 {
    color: #e8f5ff !important;
    font-weight: 700;
}
/* ---------- Main Card Container ---------- */
.gr-block, .gr-panel, .gr-group, .gr-accordion {
    background: rgba(20, 20, 20, 0.65) !important;
    border: 1px solid rgba(0, 255, 180, 0.15) !important;
    border-radius: 18px !important;
    backdrop-filter: blur(10px) !important;
    padding: 18px !important;
}
/* ---------- Tabs ---------- */
.gradio-tab {
    background: #0c0c0c !important;
}
.gradio-tab button {
    background: transparent !important;
    color: #b8c7d9 !important;
    padding: 10px 16px !important;
    border-radius: 10px !important;
    font-weight: 600;
    border: none !important;
}
.gradio-tab button:hover {
    background: rgba(0, 255, 160, 0.08) !important;
}
.gradio-tab button.selected {
    background: rgba(0, 255, 160, 0.16) !important;
    color: #00ffb4 !important;
    border: 1px solid rgba(0, 255, 160, 0.35) !important;
}
/* ---------- Textboxes ---------- */
textarea, input, .gr-textbox, .gr-textbox textarea {
    background: rgba(30, 30, 30, 0.7) !important;
    border: 1px solid rgba(0, 255, 160, 0.2) !important;
    border-radius: 14px !important;
    color: #d8f0ff !important;
    padding: 12px !important;
    font-size: 15px !important;
}
/* ---------- Scroll Boxes (Summary, MCQ, QA) ---------- */
#summary_box, #mcqs_box, #qa_box, #chat_output {
    background: rgba(20, 20, 20, 0.6) !important;
    border: 1px solid rgba(0, 255, 160, 0.25) !important;
    border-radius: 16px !important;
    padding: 14px !important;
    height: 500px;
    overflow-y: scroll;
    color: #eafffa !important;
}
/* ---------- Buttons ---------- */
button {
    background: linear-gradient(135deg, #00ffb4, #00c78c) !important;
    color: #000 !important;
    font-weight: 700 !important;
    border-radius: 12px !important;
    padding: 12px 20px !important;
    border: none !important;
    transition: 0.25s ease;
}
button:hover {
    transform: scale(1.03);
    background: linear-gradient(135deg, #00ffcf, #00e0a4) !important;
}
/* ---------- File Upload Box ---------- */
.gr-file-upload {
    background: rgba(20, 20, 20, 0.7) !important;
    border: 2px dashed rgba(0, 255, 160, 0.35) !important;
    border-radius: 16px !important;
    padding: 20px !important;
}
.gr-file-upload:hover {
    border-color: #00ffb4 !important;
}
"""


# ---------------- UI ----------------
with gr.Blocks() as ui:
    gr.HTML("<style>" + custom_css + "</style>")
    gr.Markdown("# 📚 StudyAssistant AI")

    with gr.Row():
        with gr.Column(scale=1):
            with gr.Tab("📄 Documents"):
                doc_files = gr.File(label="Upload PDF / DOCX / PPTX", file_count="multiple")
                doc_process = gr.Button("Process Documents")

            with gr.Tab("🌐 Website"):
                website_url = gr.Textbox(label="Enter Website URL")
                website_process = gr.Button("Process Website")

            with gr.Tab("💬 Chatbot"):
                chat_input = gr.Textbox(label="Ask a Question")
                chat_button = gr.Button("Ask")
                chat_output = gr.Textbox(
                    label="Answer", 
                    lines=10, 
                    interactive=True, 
                    elem_id="chat_output"
                )

        with gr.Column(scale=2):
            with gr.Tab("📑 Summary"):
                summary_box = gr.Textbox(lines=25, interactive=True, elem_id="summary_box")

            with gr.Tab("📝 MCQs"):
                mcqs_box = gr.Textbox(lines=25, interactive=True, elem_id="mcqs_box")

            with gr.Tab("❓ Short Q/A"):
                qa_box = gr.Textbox(lines=25, interactive=True, elem_id="qa_box")

# -------- Button Click Bindings (INSIDE Blocks context!) --------
    doc_process.click(process_documents, inputs=doc_files, outputs=[summary_box, mcqs_box, qa_box])
    website_process.click(process_website, inputs=website_url, outputs=[summary_box, mcqs_box, qa_box])
    chat_button.click(ask_question, inputs=chat_input, outputs=chat_output)


# -------- Launch UI --------
ui.launch()