import os
import time
import gradio as gr
import google.generativeai as genai

# 設定 Gemini API KEY
def configure_api_key(api_key):
    genai.configure(api_key=api_key)
    print("API Key 已配置成功")

# 上傳 PDF 並等待處理完成
def upload_and_process_pdf(file_path, mime_type="application/pdf"):
    print("上傳 PDF 中...")
    pdf_file = genai.upload_file(file_path, mime_type=mime_type)
    print(f"PDF '{pdf_file.display_name}' 上傳成功，URI 為: {pdf_file.uri}")
    
    # 等待 PDF 處理完成
    while pdf_file.state.name == "PROCESSING":
        print("等待 PDF 處理中...")
        time.sleep(10)
        pdf_file = genai.get_file(pdf_file.name)

    if pdf_file.state.name == "FAILED":
        raise ValueError("PDF 處理失敗。")
    
    print(f"PDF 處理完成: {pdf_file.uri}")
    return pdf_file

# 使用 PDF 的 URI 來生成描述
def generate_pdf_summary(api_key, pdf_file_path, prompt="仔細讀檔，彙整重點。繁體中文"):
    configure_api_key(api_key)

    # 上傳並處理 PDF
    try:
        pdf_file = upload_and_process_pdf(pdf_file_path)
    except Exception as e:
        return f"PDF 上傳或處理失敗：{e}", None

    # 設定模型
    model = genai.GenerativeModel(model_name="models/gemini-1.5-pro-002")

    # 發送 LLM 推理請求
    try:
        print("Making LLM inference request...")
        response = model.generate_content(
            [prompt, pdf_file],
            request_options={"timeout": 600}
        )
        # 保存模型回應到文件
        output_file_path = "/tmp/model_response.txt"
        with open(output_file_path, "w") as f:
            f.write(response.text)
        return response.text, output_file_path
    except Exception as e:
        return f"與模型對話時發生錯誤：{e}", None

# Gradio 介面
with gr.Blocks() as demo:
    gr.Markdown("### 上傳 PDF 並生成摘要（Gemini）")

    # API Key 輸入
    api_key_input = gr.Textbox(
        label="輸入 API Key",
        placeholder="請輸入您的 Gemini API Key",
        type="password"
    )

    # PDF 上傳
    pdf_input = gr.File(
        label="上傳 PDF",
        type="filepath",
    )

    # 描述提示輸入
    prompt = gr.Textbox(
        label="摘要提示",
        placeholder="默認為 '仔細讀檔，彙整重點。繁體中文'",
        lines=5
    )

    # 提交按鈕
    submit_button = gr.Button("提交")

    # 輸出文字框和下載文件按鈕
    output_text = gr.Textbox(
        label="模型回應",
        placeholder="模型的回應將出現在這裡",
        lines=10
    )
    download_button = gr.File(
        label="下載模型回應",
        interactive=False
    )

    # 設置提交按鈕的函數調用
    def handle_generate_summary(api_key, pdf_file, prompt):
        text_response, file_path = generate_pdf_summary(api_key, pdf_file, prompt)
        return text_response, file_path

    # 提交按鈕的處理
    submit_button.click(
        handle_generate_summary,
        inputs=[api_key_input, pdf_input, prompt],
        outputs=[output_text, download_button]
    )

# 啟動 Gradio 介面
demo.launch()