import gradio as gr
import subprocess
import os
import shutil

# Hàm kiểm tra mật khẩu
def check_password(password):
    if password == "thienphuc":
        return gr.update(visible=False), gr.update(visible=True), ""
    return gr.update(visible=True), gr.update(visible=False), "❌ Mã không đúng, thử lại nhé!"

def translate_pdf(pdf_file, provider, custom_api_key):
    if not pdf_file:
        return None, "Vui lòng tải lên một tệp PDF."
    
    if provider == "DeepSeek (deepseek-chat)":
        endpoint = "https://api.deepseek.com/v1"
        model_name = "deepseek-chat"
    elif provider == "Qwen (qwen-plus)":
        endpoint = "https://dashscope.aliyuncs.com/compatible-mode/v1"
        model_name = "qwen-plus"
    else:
        return None, "Vui lòng chọn mô hình hợp lệ."

    api_key = custom_api_key if custom_api_key.strip() else os.environ.get("DEFAULT_API_KEY", "")
    
    if not api_key:
        return None, "Lỗi: Không tìm thấy API Key. Vui lòng nhập API Key để tiếp tục."

    file_path = pdf_file.name
    output_dir = "output_translated"
    os.makedirs(output_dir, exist_ok=True)
    
    base_name = os.path.basename(file_path)
    name_without_ext = os.path.splitext(base_name)[0]
    expected_output_file = os.path.join(os.path.dirname(file_path), f"{name_without_ext}-vi.pdf")

    command = [
        "pdf2zh", file_path,
        "-lo", "vi",
        "-s", "openai",
        "-tk", api_key,
        "-ep", endpoint,
        "-m", model_name
    ]

    try:
        subprocess.run(command, check=True, capture_output=True, text=True)
        if os.path.exists(expected_output_file):
            final_path = os.path.join(output_dir, f"{name_without_ext}_ThienPhuc_Translated.pdf")
            shutil.move(expected_output_file, final_path)
            return final_path, "✅ Dịch thành công! Cấu trúc gốc đã được giữ nguyên."
        else:
            return None, "❌ Lỗi: Xử lý xong nhưng không tìm thấy file PDF đầu ra."
    except subprocess.CalledProcessError as e:
        return None, f"❌ Có lỗi xảy ra trong quá trình dịch:\n{e.stderr}"

# Thiết kế giao diện chính (Đã xóa tham số theme ở đây để sửa lỗi Gradio 6.0+)
with gr.Blocks(title="PDF Translator by ThienPhuc") as app:
    
    # 1. Giao diện nhập mã (Hiển thị đầu tiên)
    with gr.Group(visible=True) as login_screen:
        gr.Markdown("## 🔒 Hệ thống nội bộ")
        gr.Markdown("Vui lòng nhập mã truy cập để sử dụng hệ thống dịch PDF.")
        pwd_input = gr.Textbox(label="Mã truy cập", type="password", placeholder="Nhập mã...")
        error_msg = gr.Markdown("")
        login_btn = gr.Button("Mở khóa 🔓", variant="primary")

    # 2. Giao diện công cụ (Bị ẩn đi, chỉ hiện khi nhập đúng mã)
    with gr.Group(visible=False) as main_screen:
        gr.Markdown("# 📄 Hệ thống Dịch PDF của ThienPhuc")
        gr.Markdown("Tải tệp PDF lên, chọn bộ não AI (DeepSeek/Qwen) và nhập API Key để dịch sang Tiếng Việt mà không làm vỡ cấu trúc.")
        
        with gr.Row():
            with gr.Column(scale=1):
                file_input = gr.File(label="Tải lên tệp PDF cần dịch", file_types=[".pdf"])
                provider_dropdown = gr.Dropdown(
                    choices=["DeepSeek (deepseek-chat)", "Qwen (qwen-plus)"],
                    value="DeepSeek (deepseek-chat)",
                    label="Chọn Mô hình AI"
                )
                api_key_input = gr.Textbox(
                    label="Nhập API Key",
                    placeholder="Dán API Key của bạn (hoặc bạn bè) vào đây...",
                    type="password"
                )
                translate_btn = gr.Button("🚀 Bắt đầu dịch PDF", variant="primary")
                
            with gr.Column(scale=1):
                file_output = gr.File(label="Tải xuống bản dịch")
                status_output = gr.Textbox(label="Trạng thái hệ thống", interactive=False)

        translate_btn.click(
            fn=translate_pdf,
            inputs=[file_input, provider_dropdown, api_key_input],
            outputs=[file_output, status_output]
        )

    # Xử lý sự kiện bấm nút mở khóa
    login_btn.click(
        fn=check_password,
        inputs=[pwd_input],
        outputs=[login_screen, main_screen, error_msg]
    )

# Chuyển tham số theme xuống hàm launch
app.launch(theme=gr.themes.Soft(), ssr_mode=False)