import gradio as gr
from huggingface_hub import InferenceClient

spotify_embed = """
<iframe data-testid="embed-iframe" 
style="border-radius:12px" 
src="https://open.spotify.com/embed/playlist/3fZeYquijfHzPFtIMXzzUY?utm_source=generator" 
width="100%" 
height="352" 
frameBorder="0" 
allowfullscreen="" 
allow="autoplay; 
clipboard-write; 
encrypted-media; 
fullscreen; 
picture-in-picture" 
loading="lazy"></iframe>"""

theme = gr.themes.Soft().set(

    body_background_fill="#fff7fb",

    block_background_fill="#ffffffcc",

    border_color_primary="#f8c8dc",

    button_primary_background_fill="#ffb3c6",
    button_primary_background_fill_hover="#ff8fab",
    button_primary_text_color="#ffffff",

    button_secondary_background_fill="#e0c3fc",
    button_secondary_background_fill_hover="#cdb4db",
    button_secondary_text_color="#4a4a4a",

    body_text_color="#5c5470",
    block_title_text_color="#9d4edd",
    block_label_text_color="#7b2cbf",

    input_background_fill="#fff0f6",
    input_border_color="#f3c4d7",

    link_text_color="#c77dff"
)

with open("knowledge.txt" , "r", encoding="utf-8") as f:
    knowledge_base = f.read()

client = InferenceClient("Qwen/Qwen2.5-7B-Instruct")

SYSTEM_MESSAGES = {
    "wellness": (
        "You are a kind wellness chatbot. "
        "Give practical, supportive, and thoughtful advice "
        "about the issues the user shares."
    ),

    "story": (
        "You are a creative storytelling assistant. "
        "Create imaginative, engaging, and detailed stories "
        "based on the user's ideas."
    )
}


def respond(message, history, mode):

    if mode is None:

        user_choice = message.lower().strip()

        if user_choice in ["wellness", "wellness mode"]:
            mode = "wellness"
            yield (
                "🌿 Wellness mode activated!\n\n"
                "Tell me what's on your mind.",
                mode
            )
            return

        if user_choice in ["story", "storytelling", "creative"]:  #changed elif to if to let the user change choices anytime
            mode = "story"
            yield (
                "📖 Creative Storytelling mode activated!\n\n"
                "Give me a story idea!",
                mode
            )
            return

        else:
            yield (
                "Hi! How can I help you today?\n\n"
                "• wellness\n"
                "• story",
                mode
            )
            return

    messages = [
        {
            "role": "system",
            "content": SYSTEM_MESSAGES[mode] + "\n\n" + knowledge_base
        }
    ]

    if history:
        messages.extend(history)

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

    for chunk in client.chat_completion(
        messages=messages,
        max_tokens=1024,
        temperature=0.7,
        top_p=0.9,
        stream=True,
    ):

        token = chunk.choices[0].delta.content or "" 

        if token:
            response += token
            yield response, mode


with gr.Blocks(theme=theme) as demo:

    mode_state = gr.State(None)

    gr.Image("Website Banner.png", show_label=False, container=False)

    gr.Markdown("# ***Wellness and Storytelling ChatBot***")

    gr.Markdown(
        "DISCLAIMER: While signing, please spell out individual letters in ASL instead of the distinct signs \n\n"
        "Choose a mode by typing:\n\n"
        "- `wellness`\n"
        "- `story`"
    )

    chatbot = gr.ChatInterface(
        fn=respond,
        additional_inputs=[mode_state],
        additional_outputs=[mode_state],
    )

    gr.Markdown("### 🎵 Music We Listened To While Building This")

    gr.HTML(spotify_embed)

demo.launch(debug=True)