import openai
import gradio as gr
import fitz  # PyMuPDF

api_key = ""
selected_model = ""
summary_text = ""

def set_api_key(user_api_key):
    global api_key
    api_key = user_api_key
    return "✅ API Key 已設定"

def set_model(model_name):
    global selected_model
    selected_model = model_name
    return f"✅ 模型已選：{model_name}"

def extract_pdf_text(file_path):
    doc = fitz.open(file_path)  # 直接用檔案路徑打開
    text = ""
    for page in doc:
        text += page.get_text()
    return text

def generate_summary(pdf_file):
    global summary_text
    client = openai.OpenAI(api_key=api_key)
    pdf_text = extract_pdf_text(pdf_file)
    if not pdf_text.strip():
        return "⚠️ 無法解析 PDF 文字，可能為純圖片 PDF。"
    response = client.chat.completions.create(
        model=selected_model,
        messages=[
            {"role": "system", "content": "請將以下 PDF 內容整理為條列式摘要重點。"},
            {"role": "user", "content": pdf_text[:4000]}
        ]
    )
    summary_text = response.choices[0].message.content
    return summary_text

def ask_question(user_question):
    client = openai.OpenAI(api_key=api_key)
    response = client.chat.completions.create(
        model=selected_model,
        messages=[
            {"role": "system", "content": f"根據以下 PDF 摘要內容回答問題：\n{summary_text}"},
            {"role": "user", "content": user_question}
        ]
    )
    return response.choices[0].message.content

with gr.Blocks() as demo:
    gr.Markdown("# 📄 PDF 摘要 & 問答助手 (Hugging Face 版)")

    api_key_input = gr.Textbox(label="輸入 OpenAI API Key", type="password")
    api_key_status = gr.Textbox(label="狀態", interactive=False)
    api_key_input.submit(set_api_key, inputs=api_key_input, outputs=api_key_status)

    model_choice = gr.Radio(["gpt-4", "gpt-4.1-nano", "gpt-4o"], label="選擇模型")
    model_status = gr.Textbox(label="模型狀態", interactive=False)
    model_choice.change(set_model, inputs=model_choice, outputs=model_status)

    pdf_upload = gr.File(label="上傳 PDF")
    summary_output = gr.Textbox(label="PDF 摘要", lines=10)
    summary_btn = gr.Button("生成摘要")
    summary_btn.click(generate_summary, inputs=pdf_upload, outputs=summary_output)

    question_input = gr.Textbox(label="請輸入您的問題")
    answer_output = gr.Textbox(label="AI 回答", lines=5)
    question_btn = gr.Button("送出問題")
    question_btn.click(ask_question, inputs=question_input, outputs=answer_output)

demo.launch()