import os
import gradio as gr
from huggingface_hub import InferenceClient

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

def generate_image(prompt, model_id):
    # 安全檢查
    if not prompt:
        return None, "請輸入描述文字！"
    if not model_id:
        return None, "請輸入或選擇一個模型 ID！"
    if not HF_TOKEN:
        return None, "錯誤：請在 Space Settings > Secrets 設定 HF_TOKEN。"

    try:
        # 使用使用者輸入或選擇的模型 ID
        image = client.text_to_image(
            prompt, 
            model=model_id.strip()
        )
        return image, f"成功！使用模型：{model_id}"
    except Exception as e:
        error_msg = f"發生錯誤：{str(e)}"
        return None, error_msg

# 2. 建立介面
with gr.Blocks() as demo:
    gr.Markdown("# 🎨 AI 多模型生圖工具 (2026 版)")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="1. 描述你想看到的畫面", 
                placeholder="例如：一個穿著太空衣的貓...",
                lines=3
            )
            
            # 模型選擇區
            model_input = gr.Textbox(
                label="2. 模型 ID (可手動修改)", 
                value="black-forest-labs/FLUX.1-schnell",
                placeholder="例如: stabilityai/stable-diffusion-xl-base-1.0"
            )
            
            # 常用模型快捷鍵 (選了會自動填入上面的輸入框)
            model_presets = gr.Radio(
                choices=[
                    "black-forest-labs/FLUX.1-schnell", 
                    "stabilityai/stable-diffusion-3.5-large",
                    "digiplay/2K"
                ],
                label="常用模型快捷選擇",
                value="black-forest-labs/FLUX.1-schnell"
            )
            # 當點選 Radio 時，自動更新 Textbox 的值
            model_presets.change(fn=lambda x: x, inputs=model_presets, outputs=model_input)

            run_btn = gr.Button("立即生成 ✨", variant="primary")
            status_text = gr.Markdown()
        
        with gr.Column():
            output_img = gr.Image(label="生成結果", type="pil")

    # 按下按鈕時，同時傳入 prompt 和 model_input
    run_btn.click(
        fn=generate_image, 
        inputs=[input_text, model_input], 
        outputs=[output_img, status_text]
    )

# 3. 啟動介面
demo.launch(theme=gr.themes.Soft())