import gradio as gr
from train_and_serve import generate_video_with_audio
import threading

MAX_CPU = 16  # assume availability

chat_history = []

def chatbot_response(user_msg):
    # simple echo + context
    chat_history.append(("User", user_msg))
    resp = f"I heard: '{user_msg}'. Ask me about video generation!"
    chat_history.append(("Bot", resp))
    return resp

def handle_prompt(prompt, song_text):
    vpath, apath = generate_video_with_audio(prompt, song_text=song_text)
    return vpath, apath

with gr.Blocks() as demo:
    gr.Markdown("# 🎬 FineToon Video Chat & Generator")
    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(label="Video Prompt Text")
            song = gr.Textbox(label="Optional Song Lyrics / Voice Text")
            gen_btn = gr.Button("Generate Video")
            video_out = gr.Video(label="Generated Video")
            audio_out = gr.Audio(label="Generated Audio (Song / TTS)")
        with gr.Column(scale=1):
            chat_in = gr.Textbox(label="Chat with Assistant")
            chat_out = gr.Chatbot(label="Conversation")

    gen_btn.click(handle_prompt, inputs=[prompt, song], outputs=[video_out, audio_out])
    chat_in.submit(lambda m: chatbot_response(m), inputs=chat_in, outputs=chat_out)

    demo.queue(concurrency_count=1, max_size=4).launch()