import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
import re

# Load the multilingual LLM (FLAN-T5 base) for conversational tasks
model_name = "google/flan-t5-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer, max_length=128)

def generate_llm_response(message):
    """Generate response using FLAN-T5 with multilingual prompting"""
    if not message.strip():
        return "Please say something!"
    
    # Detect if the input is in Nepali
    is_nepali = bool(re.search(r'[\u0900-\u097F]', message))
    
    # Craft a prompt based on language detection
    if is_nepali:
        prompt = f"तपाईं एक नेपाली च्याटबोट हुनुहुन्छ। प्रयोगकर्ताले भनेको कुराको जवाफ नेपालीमा दिनुहोस्: {message}"
    else:
        prompt = f"You are a friendly chatbot that can respond in English or Nepali. Respond to the user's message: {message}"
    
    # Generate response
    response = generator(prompt, max_length=128, num_return_sequences=1, temperature=0.7)[0]['generated_text']
    
    # Post-process to ensure a complete sentence
    response = response.strip()
    if not response.endswith(('.', '!', '?')):
        response += "।" if is_nepali else "."
    
    return response

def chat_function(message, history):
    """Main chat interface function"""
    if not message.strip():
        return history, ""
    
    # Generate response
    bot_response = generate_llm_response(message)
    
    # Add to history
    history.append([message, bot_response])
    
    return history, ""

# Custom CSS
css = """
.gradio-container {
    max-width: 800px !important;
    margin: auto !important;
    background-color: #1a1a2e !important;
}
.message.user {
    background-color: #e3f2fd !important;
    border-radius: 15px !important;
    padding: 10px !important;
    color: #1e1e1e !important;
}
.message.bot {
    background-color: #d1d1d1 !important;
    border-radius: 15px !important;
    padding: 10px !important;
    color: #1e1e1e !important;
}
.chatbot .message {
    color: #1e1e1e !important;
}
.input-container {
    background: linear-gradient(90deg, #667eea 0%, #764ba2 100%) !important;
    border-radius: 25px !important;
}
.input-container input {
    color: #ffffff !important;
    background: transparent !important;
}
.gradio-chatbot {
    background-color: #16213e !important;
}
"""

# Create the Gradio interface
with gr.Blocks(css=css, title="Simple Nepali Chatbot", theme=gr.themes.Default()) as demo:
    
    gr.HTML("""
    <div style="text-align: center; padding: 20px;">
        <h1>🇳🇵 नेपाली च्याटबोट</h1>
        <h2>Simple Nepali Chatbot</h2>
        <p style="font-size: 18px;">
            <strong>नेपालीमा वा अंग्रेजीमा कुराकानी गर्नुहोस्!</strong><br>
            <em>Chat in Nepali or English!</em>
        </p>
    </div>
    """)
    
    chatbot_ui = gr.Chatbot(
        value=None,
        height=400,
        show_label=False,
        container=True,
        bubble_full_width=False,
        show_copy_button=True
    )
    
    with gr.Row():
        msg_input = gr.Textbox(
            placeholder="यहाँ लेख्नुहोस् / Type here...",
            show_label=False,
            scale=4,
            lines=1,
            container=False
        )
        send_btn = gr.Button("📤 Send", scale=1, variant="primary")
        clear_btn = gr.Button("🗑️ Clear", scale=1, variant="secondary")
    
    # Example conversations
    with gr.Row():
        gr.Examples(
            examples=[
                ["नमस्ते!"],
                ["Hello!"],
                ["तपाईंको नाम के हो?"],
                ["How are you?"],
                ["What is your name?"],
                ["कस्तो छ?"],
                ["Thank you!"],
                ["धन्यवाद!"]
            ],
            inputs=msg_input,
            label="🔄 Try these examples / यी उदाहरणहरू प्रयास गर्नुहोस्"
        )
    
    # Event handlers
    msg_input.submit(
        chat_function,
        inputs=[msg_input, chatbot_ui],
        outputs=[chatbot_ui, msg_input]
    )
    
    send_btn.click(
        chat_function,
        inputs=[msg_input, chatbot_ui],
        outputs=[chatbot_ui, msg_input]
    )
    
    clear_btn.click(
        lambda: ([], ""),
        outputs=[chatbot_ui, msg_input]
    )
    
    gr.HTML("""
    <div style="text-align: center; margin-top: 20px; padding: 20px; background: #16213e; border-radius: 10px; color: #ffffff;">
        <h3>📝 About this Chatbot</h3>
        <p>This is a simple LLM-based chatbot that responds in both Nepali and English.</p>
        <p><strong>यो एक सरल LLM-आधारित च्याटबोट हो जसले नेपाली र अंग्रेजी दुवैमा जवाफ दिन्छ।</strong></p>
        <p><em>Powered by a lightweight model - works on Hugging Face Spaces! ⚡</em></p>
    </div>
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )