import gradio as gr
import os
import spaces
from huggingface_hub import InferenceClient

# -----------------------------
# 🔑 API KEY FIXED
# -----------------------------
client = InferenceClient(
    model="meta-llama/Llama-3.1-8B-Instruct",
    token=os.getenv("API_KEY")   # 👈 FIXED NAME (recommended)
)

# -----------------------------
# SYSTEM PROMPT
# -----------------------------
SYSTEM_PROMPT = """
You are AI Study Buddy, created by Areeba Iqbal.

Rules:
- Always explain step-by-step
- Give examples
- Be clear and student-friendly
- If asked who created you: "I am AI Study Buddy, created by Areeba Iqbal."
"""

# -----------------------------
# MODE CONTROL
# -----------------------------
def build_prompt(message, mode):
    mode_prompts = {
        "📚 Study Mode": "Explain simply for students with examples.",
        "💻 Coding Mode": "Act as a senior programmer. Debug and improve code.",
        "🧮 Math Solver": "Solve step-by-step with explanation.",
        "📝 Exam Prep": "Give short exam-focused answers."
    }

    return f"""
{SYSTEM_PROMPT}

Mode: {mode_prompts.get(mode, "")}

User Question:
{message}
"""

# -----------------------------
# MAIN CHAT FUNCTION
# -----------------------------
@spaces.GPU
def get_response(message, history, mode):

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    for msg in history:
        messages.append(msg)

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

    try:
        response = client.chat_completion(
            messages=messages,
            max_tokens=1024,
            temperature=0.7
        )

        return response.choices[0].message.content

    except Exception as e:
        return f"❌ Error: {e}"


# -----------------------------
# QUICK ACTIONS
# -----------------------------
def summarize(text):
    return client.chat_completion(
        messages=[{"role": "user", "content": "Summarize: " + text}],
        max_tokens=500
    ).choices[0].message.content


def quiz(text):
    return client.chat_completion(
        messages=[{"role": "user", "content": "Generate 5 MCQs: " + text}],
        max_tokens=500
    ).choices[0].message.content


def simple(text):
    return client.chat_completion(
        messages=[{"role": "user", "content": "Explain simply: " + text}],
        max_tokens=500
    ).choices[0].message.content


def study_plan(text):
    return client.chat_completion(
        messages=[{"role": "user", "content": f"Make 7-day study plan for: {text}"}],
        max_tokens=700
    ).choices[0].message.content


# -----------------------------
# UI
# -----------------------------
css = """
.main-container {
    max-width: 900px;
    margin: auto;
}
#title { text-align:center; }
#subtitle { text-align:center; color:gray; }
#footer { text-align:center; color:gray; font-size:14px; }
"""

with gr.Blocks(
    theme=gr.themes.Soft(),
    css=css,
    title="AI Study Buddy"
) as demo:

    gr.HTML("""
    <div class="main-container">
        <h1 id="title">📚 AI Study Buddy</h1>
        <p id="subtitle">Learn smarter with AI-powered guidance</p>
    </div>
    """)

    # ---------------- MODE SELECT ----------------
    mode = gr.Radio(
        ["📚 Study Mode", "💻 Coding Mode", "🧮 Math Solver", "📝 Exam Prep"],
        value="📚 Study Mode",
        label="Select Mode"
    )

    # ---------------- CHAT ----------------
    chatbot = gr.ChatInterface(
        fn=get_response,
        additional_inputs=[mode],
        examples=[
            ["Explain recursion"],
            ["Solve quadratic equation"],
            ["What is AI?"],
            ["Debug Python code"]
        ]
    )

    # ---------------- QUICK ACTIONS ----------------
    gr.Markdown("## ⚡ Quick Actions")

    quick_input = gr.Textbox(label="Quick Input")

    with gr.Row():
        gr.Button("📖 Summarize").click(summarize, quick_input, gr.Textbox())
        gr.Button("📝 Quiz").click(quiz, quick_input, gr.Textbox())
        gr.Button("💡 Simple").click(simple, quick_input, gr.Textbox())

    # ---------------- STUDY PLAN ----------------
    gr.Markdown("## 🗓️ Study Plan Generator")

    plan_input = gr.Textbox(label="Enter Topic / Exam Detail")
    plan_output = gr.Textbox(label="Plan Output")

    gr.Button("Generate Plan").click(study_plan, plan_input, plan_output)

    # ---------------- FOOTER ----------------
    gr.HTML("""
    <div id="footer">
        Created by Areeba Iqbal
    </div>
    """)

demo.launch(server_name="0.0.0.0", server_port=7860)