import gradio as gr
from huggingface_hub import InferenceClient
import numpy as np
from PIL import Image
import tflite_runtime.interpreter as tflite

# Load TFLite model at startup
interpreter = tflite.Interpreter(model_path="model.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# Load labels
with open("labels.txt", "r") as f:
    labels = [line.strip().split(" ", 1)[-1] for line in f.readlines()]

def predict_asl(frame):
    if frame is None:
        return "–"
    img = Image.fromarray(frame).resize((224, 224))
    img_array = np.array(img, dtype=np.float32) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    interpreter.set_tensor(input_details[0]['index'], img_array)
    interpreter.invoke()
    predictions = interpreter.get_tensor(output_details[0]['index'])[0]
    best_idx = int(np.argmax(predictions))
    return labels[best_idx]

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 advice."
    ),
    "story": (
        "You are a creative storytelling assistant. "
        "Create engaging stories."
    )
}

def respond(message, history, mode):
    if mode is None:
        yield "Please select a mode first.", 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,
    ):
        delta = chunk.choices[0].delta
        token = delta.content if delta and hasattr(delta, "content") else ""
        response += token
        yield response, mode


def add_letter(letter, word):
    if letter and letter != "–":
        return word + letter
    return word

def add_space(word):
    return word + " "

def clear_word():
    return ""

def send_word(word):
    return word, ""


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

    with gr.Row():
        wellness_btn = gr.Button("🌿 Wellness", variant="primary")
        story_btn = gr.Button("📖 Story", variant="secondary")

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

    wellness_btn.click(
        fn=lambda: ([{"role": "assistant", "content": "Wellness mode"}], "wellness"),
        outputs=[chatbot.chatbot, mode_state]
    )

    story_btn.click(
        fn=lambda: ([{"role": "assistant", "content": "Story mode"}], "story"),
        outputs=[chatbot.chatbot, mode_state]
    )

    gr.Markdown("### 🤟 ASL Input")
    with gr.Group():
        gr.Markdown("*Sign a letter in front of your camera, then click Add Letter to build a word and Send to chat.*")
        with gr.Row():
            with gr.Column(scale=1):
                asl_cam = gr.Image(
                    sources=["webcam"],
                    streaming=True,
                    label="Camera",
                    mirror_webcam=True,
                    height=300
                )
            with gr.Column(scale=1):
                asl_detected = gr.Textbox(label="Detected Letter", interactive=False, value="–")
                asl_word_box = gr.Textbox(label="Current Word", interactive=False, value="")
                with gr.Row():
                    asl_add_btn = gr.Button("Add Letter", variant="primary")
                    asl_space_btn = gr.Button("Space", variant="secondary")
                    asl_clear_btn = gr.Button("Clear", variant="secondary")
                asl_send_btn = gr.Button("Send to Chat ➤", variant="primary")

    word_state = gr.State("")

    asl_cam.stream(
        fn=predict_asl,
        inputs=[asl_cam],
        outputs=[asl_detected]
    )

    asl_add_btn.click(fn=add_letter, inputs=[asl_detected, word_state], outputs=[word_state]).then(
        fn=lambda w: w, inputs=[word_state], outputs=[asl_word_box]
    )
    asl_space_btn.click(fn=add_space, inputs=[word_state], outputs=[word_state]).then(
        fn=lambda w: w, inputs=[word_state], outputs=[asl_word_box]
    )
    asl_clear_btn.click(fn=clear_word, outputs=[word_state]).then(
        fn=lambda w: w, inputs=[word_state], outputs=[asl_word_box]
    )
    asl_send_btn.click(fn=send_word, inputs=[word_state], outputs=[chatbot.textbox, word_state]).then(
        fn=lambda w: w, inputs=[word_state], outputs=[asl_word_box]
    )

    gr.Markdown("### 🎵 Music")
    gr.HTML(spotify_embed)

demo.launch(debug=True, allowed_paths=["."])