import threading
import asyncio
import uuid
import gradio as gr
import edge_tts
from deep_translator import GoogleTranslator

from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration
import torch

# ===== TTS PART =====

voice_characters = {
    "English - US": {
        "Aria": "en-US-AriaNeural",
        "Jenny": "en-US-JennyNeural",
        "Guy": "en-US-GuyNeural",
    },
    "Hindi": {
        "Swara": "hi-IN-SwaraNeural",
        "Madhur": "hi-IN-MadhurNeural"
    }
}

async def generate_tts(text, voice):
    if not text.strip():
        raise ValueError("Input text is empty")
    filename = f"output_{uuid.uuid4()}.mp3"
    communicate = edge_tts.Communicate(text, voice=voice)
    await communicate.save(filename)
    return filename

def tts_wrapper(text, language, character, translation_direction):
    try:
        original_text = text.strip()
        if not original_text:
            return "Input text is empty", None

        if translation_direction == "English to Hindi":
            text = GoogleTranslator(source='en', target='hi').translate(original_text)
        elif translation_direction == "Hindi to English":
            text = GoogleTranslator(source='hi', target='en').translate(original_text)

        voice = voice_characters.get(language, {}).get(character)
        if not voice:
            return f"Voice '{character}' not found for language '{language}'", None

        filename = asyncio.run(generate_tts(text, voice))
        return text, filename

    except Exception as e:
        return f"Error: {str(e)}", None

# ✅ FIXED: Voice character list updates properly now
def get_characters(language):
    chars = list(voice_characters.get(language, {}).keys())
    default_char = chars[0] if chars else None
    return gr.update(choices=chars, value=default_char)

# ===== CHATBOT PART USING BLENDERBOT =====

model_name = "facebook/blenderbot-400M-distill"
tokenizer = BlenderbotTokenizer.from_pretrained(model_name)
model = BlenderbotForConditionalGeneration.from_pretrained(model_name)
device = torch.device("cpu")
model.to(device)

async def generate_bot_tts(text):
    voice = voice_characters["English - US"]["Aria"]
    filename = await generate_tts(text, voice)
    return filename

def chatbot_response(history, user_message):
    if history is None:
        history = []

    history.append(("User", user_message))
    conversation_text = " ".join([msg for _, msg in history]) + " " + user_message
    inputs = tokenizer([conversation_text], return_tensors="pt").to(device)

    reply_ids = model.generate(**inputs, max_length=200)
    response = tokenizer.decode(reply_ids[0], skip_special_tokens=True)

    history.append(("Bot", response))

    chat_str = ""
    for speaker, msg in history:
        chat_str += f"{speaker}: {msg}\n"

    try:
        audio_path = asyncio.run(generate_bot_tts(response))
    except Exception as e:
        audio_path = None
        print(f"TTS generation failed: {e}")

    return history, chat_str, audio_path

# ===== GRADIO UI =====

def create_app():
    with gr.Blocks() as app:
        gr.Markdown("## 🎙️ Multi-Voice AI TTS with Translation + 🤖 BlenderBot Chatbot")

        with gr.Tab("TTS"):
            text_input = gr.Textbox(label="Enter Text", placeholder="Type something...", lines=4)
            language_dropdown = gr.Dropdown(choices=list(voice_characters.keys()), value="English - US", label="Language")
            character_dropdown = gr.Dropdown(choices=list(voice_characters["English - US"].keys()), value="Aria", label="Voice Character")
            translation_dropdown = gr.Dropdown(choices=["None", "English to Hindi", "Hindi to English"], value="None", label="Translation Direction")
            tts_button = gr.Button("🔊 Generate Voice")
            output_text = gr.Textbox(label="Translated Text or Error")
            output_audio = gr.Audio(label="Generated Voice")

            # ✅ This updates the character dropdown dynamically based on selected language
            language_dropdown.change(fn=get_characters, inputs=language_dropdown, outputs=character_dropdown)

            tts_button.click(fn=tts_wrapper,
                             inputs=[text_input, language_dropdown, character_dropdown, translation_dropdown],
                             outputs=[output_text, output_audio])

        with gr.Tab("Chatbot"):
            chat_history = gr.State([])
            user_input = gr.Textbox(label="Enter your message", lines=2, placeholder="Say something...")
            chat_display = gr.Textbox(label="Chat History", interactive=False, lines=15)
            audio_output = gr.Audio(label="Bot Voice Reply")
            send_button = gr.Button("Send")

            def respond(user_message, history):
                return chatbot_response(history, user_message)

            send_button.click(fn=respond, inputs=[user_input, chat_history], outputs=[chat_history, chat_display, audio_output])
            user_input.submit(fn=respond, inputs=[user_input, chat_history], outputs=[chat_history, chat_display, audio_output])

    return app

if __name__ == "__main__":
    app = create_app()
    app.launch(share=True)