import os
import gradio as gr
from huggingface_hub import InferenceClient

# 1. 讀取 Token 並檢查
HF_TOKEN = os.getenv("HF_TOKEN")

# 初始化客戶端 (若 Token 為 None，則嘗試匿名訪問，但生圖通常需要 Token)
client = InferenceClient(api_key=HF_TOKEN)

def generate_image(prompt):
    # 安全檢查 1: 是否輸入了文字
    if not prompt or prompt.strip() == "":
        return None, "請輸入描述文字！"
    
    # 安全檢查 2: 是否有設定 Token
    if not HF_TOKEN:
        return None, "錯誤：找不到 HF_TOKEN。請在 Space Settings > Secrets 中設定。"

    try:
        # 使用 FLUX 或指定的模型
        image = client.text_to_image(
            prompt, 
            #model="black-forest-labs/FLUX.1-schnell" # 建議用這個，比 digiplay 穩定且快
            model="digiplay/2K"
        )
        return image, "生成成功！"
    except Exception as e:
        # 捕捉詳細錯誤並顯示在介面
        error_msg = f"發生錯誤：{str(e)}"
        print(error_msg)
        return None, error_msg

# 2. 建立介面 (符合 Gradio 6.0 規範)
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 2026 AI 快速生圖工具")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="描述你想看到的畫面", 
                placeholder="例如：台北街頭的賽博龐克風格...",
                lines=3
            )
            run_btn = gr.Button("立即生成 ✨", variant="primary")
            status_text = gr.Markdown() # 用來顯示錯誤訊息或狀態
        
        with gr.Column():
            output_img = gr.Image(label="生成結果", type="pil")

    # 設定輸出包含 圖片 和 狀態文字
    run_btn.click(
        fn=generate_image, 
        inputs=input_text, 
        outputs=[output_img, status_text]
    )

# 3. 將 theme 參數放在 launch() 中
demo.launch(theme=gr.themes.Soft())