import subprocess
subprocess.run(["pip", "install", "faiss-cpu"], check=True)
import gradio as gr
import wikipedia
import numpy as np
import faiss
from gtts import gTTS
import tempfile
from langdetect import detect
import speech_recognition as sr
from pydub import AudioSegment
from transformers import pipeline
from sentence_transformers import SentenceTransformer
import os
from pydub.silence import split_on_silence
import time


# Initialize models
models = {
    'translator': pipeline('translation', model='Helsinki-NLP/opus-mt-mul-en'),
    'answer_gen': pipeline('text2text-generation', model='google/flan-t5-base'),
    'encoder': SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
}

# Add translation models
for lang in ['fr', 'ar', 'zh', 'es']:
    models[f'en_to_{lang}'] = pipeline(f'translation_en_to_{lang}',
                                     model=f'Helsinki-NLP/opus-mt-en-{lang}')

def translate(text, src, tgt):
    if src == tgt: return text
    if src != 'en': text = models['translator'](text)[0]['translation_text']
    if f'en_to_{tgt}' in models: return models[f'en_to_{tgt}'](text)[0]['translation_text']
    return text

def text_to_speech(text, lang):
    try:
        tts = gTTS(text=text, lang=lang)
        audio_path = tempfile.mktemp(suffix='.mp3')
        tts.save(audio_path)
        return audio_path
    except Exception as e:
        print(f"TTS Error: {e}")
        return None

def process_audio(audio_path):
    recognizer = sr.Recognizer()
    sound = AudioSegment.from_file(audio_path)
    chunks = split_on_silence(sound,
        min_silence_len=500,
        silence_thresh=sound.dBFS-14,
        keep_silence=500
    )

    full_text = ""
    for chunk in chunks:
        chunk_path = tempfile.mktemp(suffix='.wav')
        chunk.export(chunk_path, format="wav")
        with sr.AudioFile(chunk_path) as source:
            audio = recognizer.record(source)
            try:
                text = recognizer.recognize_google(audio)
                full_text += f" {text}"
            except:
                continue
        os.unlink(chunk_path)

    return full_text.strip() if full_text else None

def get_wikipedia_content(topic):
    try:
        wikipedia.set_lang('en')
        try:
            page = wikipedia.page(topic, auto_suggest=False)
            return page.summary[:1000]
        except wikipedia.exceptions.DisambiguationError as e:
            page = wikipedia.page(e.options[0])
            return page.summary[:1000]
    except Exception as e:
        print(f"Wikipedia error: {e}")
        return None

def generate_response(text, topic, lang):
    context = get_wikipedia_content(topic)
    if not context:
        return "Could not find information. Please try another topic.", None

    prompt = f"Context: {context}\nQuestion: {text}\nAnswer:"
    answer = models['answer_gen'](prompt, max_length=200)[0]['generated_text']
    translated = translate(answer, 'en', lang) if lang != 'en' else answer
    audio_path = text_to_speech(translated, lang)

    return translated, audio_path

def handle_interaction(audio, text, topic, lang, chat_history):
    if audio is not None:
        recognized_text = process_audio(audio)
        if recognized_text:
            text = recognized_text
        else:
            chat_history.append(("", "Could not understand audio. Please try again."))
            return chat_history, "", None

    if not text.strip():
        chat_history.append(("", "Please enter a question."))
        return chat_history, "", None

    response, audio_output = generate_response(text, topic, lang)
    chat_history.append((text, response))

    return chat_history, "", audio_output

# Custom CSS with light blue and dark blue theme
custom_css = """
.gradio-container {
    background: #f0f8ff !important;
    border: 3px solid #00008b !important;
    border-radius: 10px !important;
    font-family: 'Arial', sans-serif;
}
.gr-box {
    background-color: #e6f2ff !important;
    border: 2px solid #00008b !important;
    border-radius: 8px !important;
}
.gr-button {
    background-color: #4d94ff !important;
    border: 2px solid #00008b !important;
    color: white !important;
    border-radius: 6px !important;
}
.gr-button:hover {
    background-color: #1a75ff !important;
}
.gr-chatbot {
    background-color: #e6f2ff !important;
    border: 2px solid #00008b !important;
    border-radius: 8px !important;
}
.gr-textbox, .gr-dropdown, .gr-audio {
    background-color: #e6f2ff !important;
    border: 2px solid #00008b !important;
    border-radius: 6px !important;
}
.welcome-header {
    text-align: center;
    color: #00008b !important;
    margin-bottom: 20px;
}
.welcome-message {
    background-color: #e6f2ff;
    padding: 20px;
    border-radius: 10px;
    border: 2px solid #00008b;
    margin-bottom: 20px;
}
.avatar {
    width: 80px;
    height: 80px;
    margin: 0 auto;
    display: block;
}
"""

# Welcome page content
welcome_html = """
<div class="welcome-header">
    <img src="https://i.imgur.com/6wBs5mO.png" class="avatar" alt="AI Assistant">
    <h1>Welcome to Your Multilingual AI Assistant! 🌍</h1>
</div>
<div class="welcome-message">
    <h3>Hello! I'm your personal Wikipedia assistant 🤖</h3>
    <p>I can help you find information on any topic in multiple languages. Here's what I can do:</p>
    <ul>
        <li>🔍 Answer questions from Wikipedia knowledge</li>
        <li>🗣️ Understand both voice and text input</li>
        <li>🌐 Respond in English, French, Spanish, Chinese, or Arabic</li>
        <li>🔊 Speak answers back to you</li>
    </ul>
    <p>To get started, simply type your question or click the microphone to speak!</p>
</div>
"""

with gr.Blocks(css=custom_css, title="🌍 Multilingual AI Assistant") as demo:
    # Welcome page
    gr.HTML(welcome_html)

    # Main interface
    with gr.Row():
        with gr.Column(scale=1):
            audio_input = gr.Audio(
                sources=["microphone", "upload"],
                type="filepath",
                label="🎤 Speak or upload audio",
                interactive=True
            )
            topic_input = gr.Textbox(
                "Artificial Intelligence",
                label="📚 Wikipedia Topic"
            )
            lang_input = gr.Dropdown(
                ["en", "fr", "es", "zh", "ar"],
                value="en",
                label="🌐 Output Language"
            )

        with gr.Column(scale=2):
            chatbot = gr.Chatbot(label="Conversation")
            text_input = gr.Textbox(
                placeholder="Type your question here...",
                label="✏️ Or type here"
            )
            with gr.Row():
                clear_btn = gr.Button("🗑️ Clear Chat")
                submit_btn = gr.Button("🚀 Submit", variant="primary")

    audio_output = gr.Audio(label="🔊 Answer", visible=True)

    # Event handlers
    submit_btn.click(
        handle_interaction,
        inputs=[audio_input, text_input, topic_input, lang_input, chatbot],
        outputs=[chatbot, text_input, audio_output]
    )
    text_input.submit(
        handle_interaction,
        inputs=[audio_input, text_input, topic_input, lang_input, chatbot],
        outputs=[chatbot, text_input, audio_output]
    )
    clear_btn.click(
        lambda: ([], "", None),
        outputs=[chatbot, text_input, audio_output]
    )

demo.launch(share=True)