import gradio as gr
import os
import shutil
from docx import Document
from pptx import Presentation
import openai

def translate_document(file, target_language, target_country, api_key):
    log_messages = []
    output_path = None  # 初始化 output_path

    try:
        log_messages.append("Starting translation process...")
        
        # 複製文件到新文件
        original_file_path = file.name
        copied_file_path = "copy_" + os.path.basename(original_file_path)
        shutil.copy(original_file_path, copied_file_path)
        log_messages.append(f"Copied file to {copied_file_path}")

        openai.api_key = api_key  # 設定 OpenAI API 金鑰

        # 處理 .docx 文件
        if copied_file_path.endswith(".docx"):
            doc = Document(copied_file_path)

            # 翻譯並替換段落文本
            for paragraph in doc.paragraphs:
                if paragraph.text.strip():  # 只翻譯非空段落
                    translated_text = openai.ChatCompletion.create(
                        model="gpt-3.5-turbo",
                        messages=[{"role": "user", "content": f"Translate this to {target_language}: {paragraph.text}"}]
                    )["choices"][0]["message"]["content"]
                    paragraph.text = translated_text  # 用翻譯後的文本替換原始文本
            
            # 處理表格中的文本
            for table in doc.tables:
                for row in table.rows:
                    for cell in row.cells:
                        if cell.text.strip():  # 只翻譯非空單元格
                            translated_text = openai.ChatCompletion.create(
                                model="gpt-3.5-turbo",
                                messages=[{"role": "user", "content": f"Translate this to {target_language}: {cell.text}"}]
                            )["choices"][0]["message"]["content"]
                            cell.text = translated_text  # 用翻譯後的文本替換原始文本

            output_path = "translated_" + os.path.basename(original_file_path)
            doc.save(output_path)
            log_messages.append(f"Saved translated document as {output_path}.")

        # 處理 .pptx 文件
        elif copied_file_path.endswith(".pptx"):
            prs = Presentation(copied_file_path)

            # 翻譯並替換文本
            for slide in prs.slides:
                for shape in slide.shapes:
                    if hasattr(shape, "text") and shape.text.strip():  # 只翻譯非空文本框
                        translated_text = openai.ChatCompletion.create(
                            model="gpt-3.5-turbo",
                            messages=[{"role": "user", "content": f"Translate this to {target_language}: {shape.text}"}]
                        )["choices"][0]["message"]["content"]
                        shape.text = translated_text  # 用翻譯後的文本替換原始文本

            output_path = "translated_" + os.path.basename(original_file_path)
            prs.save(output_path)
            log_messages.append(f"Saved translated presentation as {output_path}.")

        else:
            log_messages.append("Unsupported file type. Please upload a .docx or .pptx file.")
            return None, "\n".join(log_messages)  # 返回 None 如果文件類型不支持

        return output_path, "\n".join(log_messages)  # 返回文件路徑和日誌信息

    except Exception as e:
        log_messages.append(f"An error occurred: {e}")
        return None, "\n".join(log_messages)  # 返回 None 和日誌

iface = gr.Interface(
    fn=translate_document,
    inputs=[
        gr.File(label="上傳文件 (.docx 或 .pptx)"),
        gr.Dropdown(choices=["English", "Spanish", "French", "German", "Japanese", "Korean", "Chinese", "Vietnamese"], label="Target Language"),
        gr.Dropdown(choices=["US", "UK", "ES", "FR", "DE", "JP", "KR", "TW", "VN"], label="Target Country"),
        gr.Textbox(label="API 金鑰（必填）", placeholder="請輸入您的 OpenAI API 金鑰")
    ],
    outputs=[
        gr.File(label="Download Translated Document"),
        gr.Textbox(label="Log Output", lines=10)  # 添加一個文本框來顯示日誌信息
    ],
    title="Document Translator",
    description="Translate your documents while preserving the original format.",
)

iface.launch()