import openai
import gradio as gr
import fitz  # PyMuPDF
from openai import OpenAI
import traceback

# 全域變數
api_key = ""
selected_model = "gpt-4"
summary_text = ""
client = None
pdf_text = ""

def set_api_key(user_api_key):
    """設定 OpenAI API Key 並初始化客戶端"""
    global api_key, client
    try:
        api_key = user_api_key.strip()
        if not api_key:
            return "❌ API Key 不能為空"
        
        # 支援新舊 key 格式
        if not (api_key.startswith('sk-') or api_key.startswith('sk-proj-')):
            return "❌ API Key 格式錯誤，必須以 'sk-' 或 'sk-proj-' 開頭"
        
        client = OpenAI(api_key=api_key)
        
        # 測試 API Key 是否有效
        test_response = client.chat.completions.create(
            model="gpt-4",
            messages=[{"role": "user", "content": "你好"}],
            max_tokens=5
        )
        return "✅ API Key 已設定並驗證成功！"
    except Exception as e:
        if "incorrect_api_key" in str(e).lower():
            return "❌ API Key 無效，請檢查是否正確"
        elif "quota" in str(e).lower():
            return "⚠️ API Key 有效，但配額不足"
        else:
            return f"❌ API Key 設定失敗: {str(e)}"

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

def extract_pdf_text(file_path):
    try:
        doc = fitz.open(file_path)
        text = ""
        for page_num, page in enumerate(doc):
            page_text = page.get_text()
            if page_text.strip():
                text += f"\n--- 第 {page_num + 1} 頁 ---\n{page_text}"
        doc.close()
        return text
    except Exception as e:
        return f"❌ PDF 解析錯誤: {str(e)}"

def generate_summary(pdf_file):
    global summary_text, pdf_text
    if not client:
        return "❌ 請先設定 OpenAI API Key"
    if not pdf_file:
        return "❌ 請先上傳 PDF 文件"
    try:
        pdf_text = extract_pdf_text(pdf_file.name)
        if not pdf_text.strip():
            return "⚠️ 無法解析 PDF 文字，可能為純圖片 PDF 或空白文件。"
        pdf_text_truncated = pdf_text[:8000]
        response = client.chat.completions.create(
            model=selected_model,
            messages=[
                {"role": "system", "content": "請將以下 PDF 內容整理為條列式摘要，用繁體中文回答："},
                {"role": "user", "content": pdf_text_truncated}
            ],
            temperature=0.3
        )
        summary_text = response.choices[0].message.content
        return summary_text
    except Exception as e:
        print(traceback.format_exc())
        return f"❌ 摘要生成失敗: {str(e)}"

def ask_question(user_question):
    if not client:
        return "❌ 請先設定 OpenAI API Key"
    if not summary_text and not pdf_text:
        return "❌ 請先生成 PDF 摘要"
    if not user_question.strip():
        return "❌ 請輸入問題"
    try:
        context = f"PDF 摘要:\n{summary_text}\n\n原始內容（部分）:\n{pdf_text[:2000]}"
        response = client.chat.completions.create(
            model=selected_model,
            messages=[
                {"role": "system", "content": f"根據以下 PDF 內容回答問題，請用繁體中文回答：\n{context}"},
                {"role": "user", "content": user_question}
            ],
            temperature=0.2
        )
        return response.choices[0].message.content
    except Exception as e:
        print(traceback.format_exc())
        return f"❌ 問答生成失敗: {str(e)}"

def clear_all():
    global summary_text, pdf_text
    summary_text = ""
    pdf_text = ""
    return "", "", ""

with gr.Blocks(
    title="PDF 摘要助手",
    css="""
    .gradio-container {
        max-width: none !important;
        width: 100% !important;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
        min-height: 100vh;
    }
    .main-content {
        max-width: 1600px !important;
        margin: 20px auto !important;
        padding: 30px !important;
        background: rgba(255, 255, 255, 0.95) !important;
        border-radius: 20px !important;
    }
    """
) as demo:
    with gr.Column():
        gr.Markdown("## 📄 PDF 摘要 & 問答助手")

        with gr.Tab("🔧 設定"):
            api_key_input = gr.Textbox(label="🔑 輸入 OpenAI API Key", type="password")
            api_key_status = gr.Textbox(label="API 狀態", interactive=False, value="等待設定 API Key...")
            api_key_btn = gr.Button("確認 API Key")
            api_key_btn.click(set_api_key, inputs=api_key_input, outputs=api_key_status)

            model_choice = gr.Radio(["gpt-4", "gpt-4.1", "gpt-4.5"], label="選擇 AI 模型", value="gpt-4")
            model_status = gr.Textbox(label="模型狀態", interactive=False, value="✅ 已選擇：gpt-4")
            model_choice.change(set_model, inputs=model_choice, outputs=model_status)

        with gr.Tab("📄 摘要"):
            pdf_upload = gr.File(label="上傳 PDF", file_types=[".pdf"])
            summary_btn = gr.Button("生成摘要")
            summary_output = gr.Textbox(label="PDF 摘要", lines=12)
            summary_btn.click(generate_summary, inputs=pdf_upload, outputs=summary_output)

        with gr.Tab("❓ 問答"):
            question_input = gr.Textbox(label="請輸入問題", lines=2)
            question_btn = gr.Button("送出問題")
            answer_output = gr.Textbox(label="AI 回答", lines=8)
            question_btn.click(ask_question, inputs=question_input, outputs=answer_output)
            question_input.submit(ask_question, inputs=question_input, outputs=answer_output)

        clear_btn = gr.Button("🗑️ 清除所有資料")
        clear_btn.click(clear_all, outputs=[summary_output, question_input, answer_output])

if __name__ == "__main__":
    demo.launch(
        show_error=True,
        share=True,
        server_name="0.0.0.0",
        server_port=7860
    )