# -*- coding: utf-8 -*-
"""Untitled2.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1Ny-gFLQxbToObZpy6sKVjXlw0UcQZdZU
"""

# -*- coding: utf-8 -*-
"""storyboard generator.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1m1JY_a-mgo6wXWmuLVOiSv61hojaoFjK

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
"""

import os
GROQ_API_KEY='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'
from groq import Groq
api_key = os.getenv("GROQ_API_KEY", GROQ_API_KEY)
client = Groq(api_key=api_key)
print("Groq client initialized successfully.")

"""### Generate Content with Groq

YouYou can now use the `client` object to interact with the Groq API. Here's an example of how to generate a text completion using the `chat.completions.create` method.
"""

models = client.models.list()
for model in models.data:
    print(model.id)

# Function to generate a storyboard
def generate_storyboard(scenario):
    if not scenario.strip():
        return "Please provide a scenario to generate the storyboard."

    messages = [
        {
            "role": "system",
            "content": """You are an AI storyteller. Generate a storyboard in a structured table with six scenes. For each scene you provide:
1) A scenario text describing what problem a persona is trying to resolve and by using what product or feature.
2) Storyline text for each scene, descriptive visual information and the purpose of the scene.
You must provide the output in structured format like table.
"""
        },
        {
            "role": "user",
            "content": f"Generate a 6-scene storyboard for: {scenario}"
        }
    ]

    completion = client.chat.completions.create(
        model="llama-3.1-8b-instant",
        messages=messages,
        temperature=1,
        max_tokens=1024,
        top_p=1,
        stream=False,
        stop=None,
    )

    return completion.choices[0].message.content

def chat_with_bot_stream(user_input):
    global conversation_history

    # Initialize conversation_history if it doesn't exist
    if 'conversation_history' not in globals():
        conversation_history = []

    # Defensive check: Ensure conversation_history only contains valid message dictionaries
    # If any item is not a dictionary or lacks 'role'/'content', reset it to prevent errors
    clean_history_needed = False
    for item in conversation_history:
        if not isinstance(item, dict) or 'role' not in item or 'content' not in item:
            clean_history_needed = True
            break
    if clean_history_needed:
        conversation_history = [] # Reset if malformed entries are found

    # Append the new user message
    conversation_history.append({"role": "user", "content": user_input})

    # Insert the system message only if it's the very first message in a new conversation
    if len(conversation_history) == 1: # This means only the user's message is present now
        conversation_history.insert(0, {
            "role": "system",
            "content": "You are an expert in storyboarding. Provide structured and insightful responses to queries about creating and refining storyboards."
        })

    completion = client.chat.completions.create(
        model="llama-3.1-8b-instant",
        messages=conversation_history,
        temperature=1,
        max_tokens=1024,
        top_p=1,
        stream=True,
        stop=None,
    )

    response_content = ""
    for chunk in completion:
        response_content += chunk.choices[0].delta.content or ""

    # Append the assistant's response
    conversation_history.append({"role": "assistant", "content": response_content})

    # The gr.Chatbot with type='messages' expects a list of dictionaries
    # directly representing the conversation history. Return the cleaned history.
    return conversation_history

import gradio as gr

TITLE = ""
CSS = """
h1 { text-align: center; font-size: 24px; margin-bottom: 10px; }
"""

with gr.Blocks(theme=gr.themes.Glass(primary_hue="violet", secondary_hue="violet", neutral_hue="stone"), css=CSS) as demo:
    with gr.Tabs():
        with gr.TabItem("Chat"):
            gr.HTML(TITLE)
            chatbot = gr.Chatbot(label="Storyboard Chatbot", type='messages', allow_tags=False)
            with gr.Row():
                user_input = gr.Textbox(
                    label="Your Message",
                    placeholder="Type your question here...",
                    lines=1
                )
                send_button = gr.Button("Ask Question")

            # Chatbot functionality
            send_button.click(
                fn=chat_with_bot_stream,
                inputs=user_input,
                outputs=chatbot,
                queue=True
            ).then(
                fn=lambda: "",
                inputs=None,
                outputs=user_input
            )

        with gr.TabItem("Generate Storyboard"):
            gr.Markdown("# Generate a Storyboard")
            scenario_input = gr.Textbox(label="Enter your scenario")
            generate_btn = gr.Button("Generate Storyboard")
            storyboard_output = gr.Textbox(label="Generated Storyboard", interactive=False)
            generate_btn.click(generate_storyboard, inputs=scenario_input, outputs=storyboard_output)

demo.launch()

get_ipython().system('pip install groq')