import gradio as gr
import os
from groq import Groq

# Initialize Groq client
client = Groq(api_key=os.getenv("GROQ_API_KEY"))

# Conversation memory
conversation_history = [
    {
        "role": "system",
        "content": "You are an expert in storyboarding. Provide structured and insightful responses."
    }
]

def chat_with_bot(user_input, history):
    history = history or conversation_history
    history.append({"role": "user", "content": user_input})

    completion = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=history,
        temperature=1,
        max_tokens=1024,
        top_p=1,
        stream=False,
    )

    assistant_msg = completion.choices[0].message.content
    history.append({"role": "assistant", "content": assistant_msg})

    return history, history

def generate_storyboard(scenario):
    if not scenario.strip():
        return "Please provide a scenario."

    messages = [
        {"role": "system", "content": "Generate a 6-scene storyboard in table form."},
        {"role": "user", "content": scenario}
    ]

    completion = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=messages,
        temperature=1,
        max_tokens=1024,
    )

    return completion.choices[0].message.content


TITLE = """
<h1 style="text-align:center;">📖 Storyboard Assistant</h1>
"""

with gr.Blocks(theme=gr.themes.Glass()) as demo:
    gr.HTML(TITLE)

    with gr.Tabs():
        with gr.Tab("💬 Chat"):
            chatbot = gr.Chatbot(type="messages", label="Storyboard Chatbot")
            state = gr.State([])

            user_input = gr.Textbox(
                placeholder="Ask about storyboards...",
                label="Your Message"
            )
            send = gr.Button("Ask")

            send.click(
                chat_with_bot,
                inputs=[user_input, state],
                outputs=[chatbot, state]
            )

        with gr.Tab("📖 Generate Storyboard"):
            scenario = gr.Textbox(label="Scenario")
            generate = gr.Button("Generate")
            output = gr.Textbox(label="Storyboard", lines=12)

            generate.click(generate_storyboard, scenario, output)

demo.launch()