import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import spaces  # <--- ĐIỀU KIỆN 1: Import thư viện này
import os

# 1. Cấu hình Model
REPO_ID = "mradermacher/Floppa-12B-Gemma3-Uncensored-i1-GGUF"
FILENAME = "Floppa-12B-Gemma3-Uncensored.i1-IQ1_S.gguf" # Nhớ thay tên file đúng

# 2. Tải model về (Bước này chỉ tải file, không tốn GPU nên để ngoài OK)
print(f"Đang tải {FILENAME}...")
try:
    model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
except Exception as e:
    # Fallback nếu tên file sai, bạn nên check lại trên HF
    print(f"Lỗi tải file: {e}")
    model_path = ""

# Biến toàn cục để lưu model
llm = None

# Hàm load model an toàn
def load_model_if_needed():
    global llm
    if llm is None:
        print("Đang load model vào GPU...")
        # n_gpu_layers=-1 để đẩy hết vào GPU
        llm = Llama(
            model_path=model_path,
            n_gpu_layers=-1, 
            n_ctx=4096,
            verbose=False
        )
    return llm

# 3. Hàm xử lý dịch thuật (CÓ GPU)
@spaces.GPU  # <--- ĐIỀU KIỆN 2: Decorator bắt buộc để kích hoạt ZeroGPU
def translate_text(text, target_lang, temperature, max_tokens):
    # Load model NGAY TRONG KHI ĐANG CÓ GPU
    # Vì ZeroGPU có thể di chuyển session, nên đảm bảo model nằm đúng chỗ
    current_llm = load_model_if_needed()
    
    prompt = f"""<start_of_turn>user
Translate the following text into {target_lang} accurately without explanation.

Text:
{text}<end_of_turn>
<start_of_turn>model
Translation:"""

    output = current_llm(
        prompt,
        max_tokens=max_tokens,
        temperature=temperature,
        stop=["<end_of_turn>", "Text:"], 
        echo=False
    )
    
    return output['choices'][0]['text'].strip()

# 4. Giao diện Gradio (Giữ nguyên)
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"# 🤖 Floppa Translator (ZeroGPU Enabled)")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Văn bản nguồn", lines=5)
            target_lang = gr.Dropdown(["Vietnamese", "English", "Japanese"], value="Vietnamese", label="Dịch sang")
            btn = gr.Button("Dịch", variant="primary")
        with gr.Column():
            output_text = gr.Textbox(label="Kết quả", lines=5)

    btn.click(translate_text, inputs=[input_text, target_lang], outputs=output_text)

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