import os
import gradio as gr
from groq import Groq
from pypdf import PdfReader

MODEL = "llama-3.1-8b-instant"
api_key = os.getenv("GROQ_API_KEY")
if not api_key:
    raise RuntimeError("GROQ_API_KEY not found")

client = Groq(api_key=api_key)
DEFAULT_MSG = "How can I help you today?"

def read_pdf(file):
    if file is None:
        return ""

    # HF sends dict: {'name':..., 'path':...}
    if isinstance(file, dict):
        path = file.get("path")
    else:
        path = getattr(file, "name", file)

    reader = PdfReader(path)
    text = ""
    for page in reader.pages:
        text += page.extract_text() or ""
    return text

def ask_ai(prompt):
    res = client.chat.completions.create(
        model=MODEL,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.4,
        max_tokens=700,
    )
    return res.choices[0].message.content

def summarizer(text): return ask_ai("Summarize:\n"+text)
def explainer(text): return ask_ai("Explain simply:\n"+text)
def mcq_generator(text): return ask_ai("Generate 10 MCQs:\n"+text)
def flashcard_generator(text): return ask_ai("Create 10 flashcards:\n"+text)

def router(text, mode):
    return {
        "Summarize": summarizer,
        "Explain": explainer,
        "MCQs": mcq_generator,
        "Flashcards": flashcard_generator,
    }[mode](text)

def app(text, pdf, mode):
    if pdf is not None:
        text = read_pdf(pdf)

    if not text.strip():
        return DEFAULT_MSG

    return router(text, mode)

custom_css = """
/* FORCE OVERRIDE HF THEME */
.gradio-container * { box-sizing: border-box; }

/* Remove radio circles */
input[type="radio"]{display:none !important;}

/* Pills */
label:has(input[type="radio"]) {
    padding:6px 12px;
    margin:4px;
    border-radius:999px;
    border:1px solid #334155;
    color:#cbd5f5;
    cursor:pointer;
}
label:has(input[type="radio"]:checked) {
    background:linear-gradient(90deg,#38bdf8,#a78bfa);
    color:black !important;
    font-weight:700;
    border:none;
}

/* Center fake message */
.center-fake textarea{
    text-align:center !important;
    font-size:22px !important;
    padding-top:40% !important;
}
"""

with gr.Blocks(css=custom_css) as interface:
    gr.Markdown("## 📘 AI Study Assistant")

    output = gr.Textbox(
        lines=10,
        label="Output",
        value=DEFAULT_MSG,
        interactive=False,
        elem_classes=["center-fake"]
    )

    text = gr.Textbox(lines=4, label="Enter Notes or Topic")
    pdf = gr.File(label="Upload PDF (optional)")
    mode = gr.Radio(
        ["Summarize","Explain","MCQs","Flashcards"],
        value="Summarize",
        label="Choose Mode"
    )
    run = gr.Button("⚡ Generate")

    def run_app(t,p,m):
        res = app(t,p,m)
        if res == DEFAULT_MSG:
            return gr.update(value=res, interactive=False, elem_classes=["center-fake"])
        else:
            return gr.update(value=res, interactive=True, elem_classes=[])

    run.click(run_app,[text,pdf,mode],output)
    text.submit(run_app,[text,pdf,mode],output)

interface.launch()