import gradio as gr
import openai
import fitz  # PyMuPDF
import os
import time


# ✅ 使用環境變數來安全存取 OpenAI API Key
openai_key = os.getenv("OPENAI_API_KEY")
if not openai_key:
    raise ValueError("API Key 未設置，請確保已設定環境變數 OPENAI_API_KEY")

# ✅ PDF 檔案名稱（將 PDF 上傳到 Space 目錄）
PDF_FILE = "statistics.pdf"

# ✅ 萃取 PDF 內容
def extract_text_from_pdf(pdf_path):
    try:
        doc = fitz.open(pdf_path)
        text = ""
        for page in doc:
            text += page.get_text()
        print(f"✅ 成功讀取 {pdf_path}")
        return text
    except Exception as e:
        print(f"❌ PDF 解析錯誤: {e}")
        return ""

# ✅ 嘗試載入 PDF 內容
if os.path.exists(PDF_FILE):
    content = extract_text_from_pdf(PDF_FILE)
else:
    print(f"⚠️ 找不到 {PDF_FILE}，請將 PDF 上傳到 Space。")
    content = ""

# ✅ 調用 OpenAI API
def openai_api(messages, openai_key):
    try:
        client = openai.OpenAI(api_key=openai_key)
        completion = client.chat.completions.create(
            model="gpt-4o",
            messages=messages
        )
        if not completion or not completion.choices:
            return "API 沒有回應，請檢查 API Key 或伺服器狀態。"
        response = completion.choices[0].message.content
        return response
    except Exception as e:
        return f"API 呼叫發生錯誤：{str(e)}"

# ✅ 準備對話訊息
def predict(inputs, chatbot):
    messages = []
    system_prompt = {
        "role": "system",
        "content": f"請扮演助教機器人，針對我所上傳的『統計學』PDF 文件進行問答。以下是學習內容：\n\n{content}"
    }
    messages.append(system_prompt)

    if chatbot is None:
        chatbot = []

    for conv in chatbot:
        if isinstance(conv, dict) and "role" in conv and "content" in conv:
            messages.append({"role": conv["role"], "content": conv["content"]})

    messages.append({"role": "user", "content": inputs})
    return messages

# ✅ 逐字輸出訊息
def slow_echo(inputs, chatbot):
    messages = predict(inputs, chatbot)
    re_message = openai_api(messages, openai_key)

    if not re_message:
        re_message = "無法取得回應，請稍後再試。"

    for i in range(len(re_message)):
        yield re_message[: i + 1]
        time.sleep(0.05)

# ✅ 建立 Gradio 介面
def setup_gradio_interface():
    demo = gr.ChatInterface(
        slow_echo,
        chatbot=gr.Chatbot(height=500, type="messages"),  # ✅ 修正 type 參數
        title="📊 統計學助教機器人",
        description="請輸入與統計學有關的問題，機器人將基於所上傳的 PDF 內容來回答。"
    )
    return demo

# ✅ 啟動應用程式
if __name__ == "__main__":
    demo = setup_gradio_interface()
    port = int(os.environ.get("PORT", 7860))
    demo.queue()
    #demo.launch(server_name="0.0.0.0", server_port=port)
    demo.launch()