import gradio as gr
import logging, langdetect, os
from huggingface_hub import InferenceClient
from deep_translator import GoogleTranslator

hf_token = os.getenv("HF_AUTH_TOKEN")

# Khởi tạo client HF và translator
client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2", token=hf_token)
translator_vi2en = GoogleTranslator(source='vi', target='en')
translator_en2vi = GoogleTranslator(source='en', target='vi')

logging.basicConfig(level=logging.INFO)

def respond(
    message,
    history: list[tuple[str, str]],
    system_message="""
    You are a professional speaking AI assistant.
    You can help solving the problems.
    You can only speak English with correct spelling, punctuation and grammar.
    Do not expose your name at all. When greeting like "hello", "hi",... just greet back without using the word "Greetings".
    """,
    max_tokens=2048,
    temperature=0.0,
    top_p=0.9,
):
    # Tạo list tin nhắn bắt đầu từ hệ thống
    messages = [{"role": "system", "content": system_message}]
    
    # Dịch câu hỏi của người dùng từ tiếng Việt sang tiếng Anh
    message_en = translator_vi2en.translate(message)
    
    # Thêm các tin nhắn lịch sử vào messages
    for val in history:
        if val[0]:  # Tin nhắn của người dùng
            messages.append({"role": "user", "content": val[0]})
        if val[1]:  # Tin nhắn của trợ lý
            messages.append({"role": "assistant", "content": val[1]})

    # Thêm câu hỏi của người dùng vào cuối lịch sử
    messages.append({"role": "user", "content": message_en})

    response = ""

    # Gửi yêu cầu tới mô hình và nhận kết quả
    try:
        for chunk in client.chat_completion(
            model="HuggingFaceH4/zephyr-7b-beta",
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            stream=True
        ):
            # Tiến hành nhận phản hồi từ mô hình
            response += chunk.choices[0].delta.content

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

    response_vi = translator_en2vi.translate(response)
    logging.info(f"Successfully generated text: {response}")
    
    return response if langdetect.detect(response) == 'vi' else response_vi

# Gradio UI
demo = gr.ChatInterface(fn=respond, theme="soft", title="Chatbot tiếng Việt", description="Chatbot hỗ trợ bài tập")

if __name__ == "__main__":
    demo.launch()