import gradio as gr
from groq import Groq
import os

# -------------------------
# Groq Client
# -------------------------
API_KEY = os.getenv("GROQ_API_KEY")

if not API_KEY:
    raise RuntimeError("❌ Please set GROQ_API_KEY in your environment variables.")

client = Groq(api_key=API_KEY)

SYSTEM_PROMPT = """You are an expert in storyboarding.
Provide structured, creative, and insightful responses about creating and refining storyboards.
Use clear sections, shots, and visual descriptions when possible.
"""

# -------------------------
# Chat Function (UPDATED)
# -------------------------
def respond(message, history, model, temperature, max_tokens):

    if message is None:
        message = ""

    message = message.strip()

    if message == "":
        return history

    # Start with system prompt
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    # Append previous conversation (NEW FORMAT)
    for msg in history:
        messages.append(msg)

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

    try:
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=temperature,
            max_completion_tokens=max_tokens,
        )

        assistant_reply = response.choices[0].message.content

        # Update history in NEW format
        history.append({"role": "user", "content": message})
        history.append({"role": "assistant", "content": assistant_reply})

        return history

    except Exception as e:
        history.append({"role": "assistant", "content": f"❌ Error: {str(e)}"})
        return history


# -------------------------
# Custom CSS
# -------------------------
custom_css = """
body {
    background: linear-gradient(135deg, #ffd966, #1e293b);
}
#title {
    text-align: center;
    font-size: 40px;
    font-weight: bold;
    color: #990000;
}
#subtitle {
    text-align: center;
    font-size: 24px;
    color: #660000;
}
"""

# -------------------------
# UI
# -------------------------
with gr.Blocks(css=custom_css) as demo:

    gr.Markdown(" 🎨 ⚡ The Magical Storyboard ⚡ 🎬 ^O^", elem_id="title")
    gr.Markdown(
        "Turn your imagination into **cinematic storyboards** using AI — one prompt at a time.",
        elem_id="subtitle",
    )

    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(height=420)
            msg = gr.Textbox(
                placeholder="Describe your scene, commercial, or movie moment...",
                label="Your Idea"
            )

            send = gr.Button("🎥 Create Storyboard")
            clear = gr.Button("🧹 Clear Chat")

        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Controls")

            model = gr.Dropdown(
                choices=[
                    "llama-3.3-70b-versatile",
                    "llama-3.1-8b-instant",
                ],
                value="llama-3.3-70b-versatile",
                label="Model",
            )

            temperature = gr.Slider(
                0, 2, value=0.9, step=0.1,
                label="Creativity (Temperature)"
            )

            max_tokens = gr.Slider(
                256, 8192, value=2048, step=256,
                label="Max Tokens"
            )

            gr.Markdown("### ✨ Examples")
            ex1 = gr.Button("☕ 30s Coffee Commercial")
            ex2 = gr.Button("👻 Horror Movie Opening")
            ex3 = gr.Button("📚 Romantic Bookstore Meet-Cute")

    # -------------------------
    # Interactions
    # -------------------------
    send.click(
        respond,
        inputs=[msg, chatbot, model, temperature, max_tokens],
        outputs=chatbot,
    )

    msg.submit(
        respond,
        inputs=[msg, chatbot, model, temperature, max_tokens],
        outputs=chatbot,
    )

    clear.click(lambda: [], None, chatbot)

    ex1.click(lambda: "Create a storyboard for a 30-second coffee commercial", None, msg)
    ex2.click(lambda: "Generate a horror movie opening scene storyboard", None, msg)
    ex3.click(lambda: "Design a storyboard for a romantic comedy meet-cute at a bookstore", None, msg)

# -------------------------
# Launch
# -------------------------
if __name__ == "__main__":
    demo.queue().launch()