import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import os

# Đường dẫn mô hình
MODEL_PATH = "content/envit5-translation"

# Hàm kiểm tra và tải mô hình nếu cần
def check_and_download_model():
    if not os.path.exists(MODEL_PATH):
        model_name = "VietAI/envit5-translation"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        
        # Lưu mô hình
        tokenizer.save_pretrained(MODEL_PATH)
        model.save_pretrained(MODEL_PATH)

# Hàm dịch
def translate_text(input_text, direction):
    if not input_text.strip():
        return "Vui lòng nhập văn bản cần dịch!"
    
    try:
        # Tải mô hình và tokenizer
        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
        model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH)
        
        # Chuẩn bị đầu vào với định hướng dịch
        prefix = "vi: " if direction == "Vie → Eng" else "en: "
        inputs = [f"{prefix}{input_text}"]
        encoded_inputs = tokenizer(inputs, return_tensors="pt", padding=True).input_ids
        
        # Dịch văn bản
        outputs = model.generate(encoded_inputs, max_length=512)
        translated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
        return translated_text
    except Exception as e:
        return f"Lỗi: {str(e)}"

# Hàm tạo file tải xuống
def create_txt_file(output_text):
    file_path = "translation_output.txt"
    with open(file_path, "w", encoding="utf-8") as f:
        f.write(output_text)
    return file_path

# Tải mô hình khi khởi chạy
check_and_download_model()

# Tạo giao diện Gradio
with gr.Blocks() as interface:
    gr.Markdown("## Ứng dụng Dịch Máy VietAI")
    gr.Markdown("Công cụ dịch tự động từ tiếng Việt sang tiếng Anh hoặc ngược lại, hỗ trợ tải kết quả dưới dạng file .txt.")
    
    # Các thành phần giao diện
    with gr.Row():
        input_text = gr.Textbox(lines=10, placeholder="Nhập văn bản cần dịch...")
        output_text = gr.Textbox(lines=10, placeholder="Kết quả bản dịch...")
    
    direction = gr.Radio(
        ["Vie → Eng", "Eng → Vie"], 
        label="Hướng dịch", 
        value="Vie → Eng"
    )
    
    with gr.Row():
        translate_button = gr.Button("Dịch")
        download_button = gr.Button("Tải xuống .txt")
    
    # Kết nối các thành phần với logic
    translate_button.click(translate_text, inputs=[input_text, direction], outputs=output_text)
    download_button.click(create_txt_file, inputs=output_text, outputs=gr.File(label="Tải file kết quả"))

# Chạy ứng dụng
if __name__ == "__main__":
    interface.launch()