import gradio as gr
from groq import Groq
import os
from fpdf import FPDF

# =========================
# INIT
# =========================
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
SYSTEM_PROMPT = "You are a helpful AI mentor."

# =========================
# ROADMAP GENERATOR
# =========================
def generate(domain, level, hours, months):
    prompt = f"""
Create a structured learning roadmap.

Domain: {domain}
Level: {level}
Hours per day: {hours}
Months: {months}

Make weekly roadmap with projects.
"""

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

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


# =========================
# PDF DOWNLOAD
# =========================
def download_pdf(text):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)

    if not text:
        text = "No roadmap generated."

    clean = text.encode("latin-1", "ignore").decode("latin-1")

    for line in clean.split("\n"):
        pdf.multi_cell(0, 8, line)

    path = "roadmap.pdf"
    pdf.output(path)
    return path


# =========================
# CHATBOT (METADATA SAFE)
# =========================
def chat(message, history):
    history = history or []

    # ---- Build messages for Groq (NO metadata)
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    for item in history:
        if isinstance(item, dict):
            role = item.get("role")
            content = item.get("content")
            if role and content:
                messages.append({"role": role, "content": content})

    messages.append({"role": "user", "content": message})

    # ---- Add to UI history
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": ""})

    try:
        stream = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=messages,
            stream=True,
        )

        reply = ""
        for chunk in stream:
            if chunk.choices[0].delta.content:
                token = chunk.choices[0].delta.content
                reply += token
                history[-1]["content"] = reply
                yield "", history

    except Exception as e:
        history[-1]["content"] = f"Error: {str(e)}"
        yield "", history


# =========================
# UI
# =========================
with gr.Blocks(title="AI Roadmap Startup") as app:

    gr.Markdown("# 🚀 AI Roadmap Generator + AI Mentor")

    with gr.Tabs():

        # -------- ROADMAP TAB
        with gr.Tab("Roadmap Generator"):
            domain = gr.Textbox(label="Domain", placeholder="Data Science")
            level = gr.Dropdown(
                ["Beginner", "Intermediate", "Advanced"],
                value="Beginner",
                label="Level",
            )
            hours = gr.Slider(1, 10, value=2, label="Hours/day")
            months = gr.Slider(1, 12, value=3, label="Months")

            gen_btn = gr.Button("Generate Roadmap")
            roadmap_output = gr.Textbox(lines=18)

            pdf_btn = gr.Button("Download PDF")
            pdf_file = gr.File()

            gen_btn.click(generate, [domain, level, hours, months], roadmap_output)
            pdf_btn.click(download_pdf, roadmap_output, pdf_file)

        # -------- CHAT TAB
        with gr.Tab("AI Chatbot"):
            chatbot = gr.Chatbot(height=500)
            msg = gr.Textbox(placeholder="Ask anything...")

            msg.submit(chat, [msg, chatbot], [msg, chatbot])

# =========================
if __name__ == "__main__":
    app.launch(theme=gr.themes.Soft())