import os
import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image
import io
import base64
import time
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional

# ============================================
# 配置
# ============================================

HF_TOKEN = os.getenv("HF_TOKEN", "")
PUBLIC_API_KEY = os.getenv("PUBLIC_API_KEY", "demo-key-123456")
SPACE_URL = os.getenv("SPACE_URL", "")

# 支援的模型
MODELS = {
    "FLUX.1-schnell (4GB, 推薦)": "black-forest-labs/FLUX.1-schnell",
    "Z-Image (8GB, 中文支援)": "Tongyi-MAI/Z-Image",
    "Hunyuan-3.0 (9GB)": "tencent/HunyuanImage-3.0-Instruct",
}

# ============================================
# FastAPI 初始化
# ============================================

app = FastAPI()

class GenerateRequest(BaseModel):
    prompt: str
    model: str = "black-forest-labs/FLUX.1-schnell"
    negative_prompt: str = ""
    guidance_scale: float = 7.5
    num_steps: int = 8
    api_key: str

class GenerateResponse(BaseModel):
    status: int
    image_base64: Optional[str] = None
    model: Optional[str] = None
    prompt: Optional[str] = None
    generation_time: Optional[float] = None
    error: Optional[str] = None

# ============================================
# API 端點
# ============================================

@app.post("/api/generate", response_model=GenerateResponse)
async def api_generate(request: GenerateRequest):
    """REST API 端點"""
    
    # API Key 驗證
    if request.api_key != PUBLIC_API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API Key")
    
    if not request.prompt or len(request.prompt) < 3:
        raise HTTPException(status_code=400, detail="Prompt too short")
    
    start_time = time.time()
    
    try:
        client = InferenceClient(token=HF_TOKEN if HF_TOKEN else None)
        
        image = client.text_to_image(
            prompt=request.prompt,
            model=request.model,
            negative_prompt=request.negative_prompt or None,
            guidance_scale=request.guidance_scale,
            num_inference_steps=request.num_steps,
        )
        
        # 轉 Base64
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        img_base64 = base64.b64encode(buffer.getvalue()).decode()
        
        return GenerateResponse(
            status=200,
            image_base64=f"data:image/png;base64,{img_base64}",
            model=request.model,
            prompt=request.prompt,
            generation_time=round(time.time() - start_time, 2),
        )
        
    except Exception as e:
        error_msg = str(e)
        if "10GB" in error_msg:
            raise HTTPException(status_code=403, detail="Model exceeds 10GB free limit")
        raise HTTPException(status_code=500, detail=error_msg)

@app.get("/health")
async def health():
    return {
        "status": "ok",
        "models": list(MODELS.values()),
        "api_key_required": True,
    }

# ============================================
# Gradio UI
# ============================================

def generate_image(prompt, model_name, negative, guidance, steps, api_key):
    """UI 生成函式"""
    
    if api_key != PUBLIC_API_KEY:
        return None, "❌ API Key 無效"
    
    if not prompt or len(prompt) < 3:
        return None, "❌ 提示詞太短（至少3個字）"
    
    model_id = MODELS.get(model_name, "black-forest-labs/FLUX.1-schnell")
    start_time = time.time()
    
    try:
        client = InferenceClient(token=HF_TOKEN if HF_TOKEN else None)
        
        image = client.text_to_image(
            prompt=prompt,
            model=model_id,
            negative_prompt=negative if negative else None,
            guidance_scale=guidance,
            num_inference_steps=steps,
        )
        
        elapsed = round(time.time() - start_time, 2)
        return image, f"✅ 生成成功！耗時 {elapsed} 秒"
        
    except Exception as e:
        error = str(e)
        if "10GB" in error:
            return None, "❌ 模型超過 10GB 免費限制，請選擇 FLUX.1-schnell"
        elif "401" in error:
            return None, "❌ HF Token 無效，請檢查 Settings > Secrets"
        return None, f"❌ 錯誤：{error}"

# 建立 Gradio 介面
with gr.Blocks(title="AI 圖片生成 API") as demo:
    
    gr.Markdown(f"""
    # 🖼️ AI 圖片生成 API + 試玩站
    
    | 項目 | 內容 |
    |------|------|
    | **API URL** | `{SPACE_URL or '請設定 SPACE_URL 環境變數'}/api/generate` |
    | **Demo API Key** | `{PUBLIC_API_KEY}` |
    | **Rate Limit** | 300 requests/hour |
    """)
    
    with gr.Tabs():
        
        # === Tab 1: 試玩 UI ===
        with gr.Tab("🎨 試玩生成"):
            with gr.Row():
                with gr.Column(scale=1):
                    prompt_input = gr.Textbox(
                        label="提示詞 (Prompt)",
                        placeholder="描述你想要的圖片，例如：A cute cat in space suit",
                        lines=3,
                    )
                    negative_input = gr.Textbox(
                        label="負面提示詞 (Negative Prompt)",
                        placeholder="不想要在圖片中的內容...",
                        lines=2,
                    )
                    model_dropdown = gr.Dropdown(
                        choices=list(MODELS.keys()),
                        value="FLUX.1-schnell (4GB, 推薦)",
                        label="選擇模型",
                    )
                    
                    with gr.Row():
                        guidance_slider = gr.Slider(
                            minimum=1.0,
                            maximum=20.0,
                            value=7.5,
                            step=0.5,
                            label="Guidance Scale (引導強度)",
                        )
                        steps_slider = gr.Slider(
                            minimum=4,
                            maximum=50,
                            value=8,
                            step=1,
                            label="Inference Steps (推理步數)",
                        )
                    
                    api_key_input = gr.Textbox(
                        label="API Key",
                        value=PUBLIC_API_KEY,
                        type="password",
                    )
                    
                    generate_btn = gr.Button("🚀 生成圖片", variant="primary")
                
                with gr.Column(scale=1):
                    output_image = gr.Image(label="生成結果", type="pil")
                    output_message = gr.Textbox(label="狀態訊息", interactive=False)
            
            # 範例
            gr.Examples(
                examples=[
                    ["A cute cat wearing astronaut helmet, digital art, colorful", "FLUX.1-schnell (4GB, 推薦)"],
                    ["夕陽下的海邊，傳統中國水墨畫風格，寧靜優美", "Z-Image (8GB, 中文支援)"],
                    ["Cyberpunk city at night, neon lights, 8k, highly detailed", "FLUX.1-schnell (4GB, 推薦)"],
                    ["Abstract geometric patterns, vibrant colors, modern art", "FLUX.1-schnell (4GB, 推薦)"],
                ],
                inputs=[prompt_input, model_dropdown],
            )
            
            generate_btn.click(
                fn=generate_image,
                inputs=[prompt_input, model_dropdown, negative_input, guidance_slider, steps_slider, api_key_input],
                outputs=[output_image, output_message],
            )
        
        # === Tab 2: API 文件 ===
        with gr.Tab("📚 API 文件"):
            gr.Markdown(f"""
            ### 可用模型
            
            | 模型 | 大小 | 免費可用 | 特點 |
            |------|------|----------|------|
            | FLUX.1-schnell | 4GB | ✅ | 4-8步快速生成，最適合 API |
            | Z-Image | 8GB | ✅ | 阿里巴巴，中文提示詞友善 |
            | Hunyuan-3.0 | 9GB | ✅ | 騰訊混元，圖像品質優秀 |
            
            ### REST API 端點
            
            **URL**: `POST {SPACE_URL or "https://your-space.hf.space"}/api/generate`
            
            **Headers**:
            ```
            Content-Type: application/json
            ```
            
            **Request Body**:
            ```json
            {{
                "prompt": "A beautiful sunset over ocean",
                "model": "black-forest-labs/FLUX.1-schnell",
                "negative_prompt": "",
                "guidance_scale": 7.5,
                "num_steps": 8,
                "api_key": "{PUBLIC_API_KEY}"
            }}
            ```
            
            **Response**:
            ```json
            {{
                "status": 200,
                "image_base64": "data:image/png;base64,iVBORw0KGgo...",
                "model": "black-forest-labs/FLUX.1-schnell",
                "prompt": "A beautiful sunset over ocean",
                "generation_time": 5.23
            }}
            ```
            """)
            
            with gr.Accordion("📜 cURL 範例", open=False):
                gr.Code(f'''curl -X POST {SPACE_URL or "https://your-space.hf.space"}/api/generate \\
  -H "Content-Type: application/json" \\
  -d '{{
    "prompt": "A cute cat in space",
    "model": "black-forest-labs/FLUX.1-schnell",
    "negative_prompt": "",
    "guidance_scale": 7.5,
    "num_steps": 8,
    "api_key": "{PUBLIC_API_KEY}"
  }}' ''', language="bash")
            
            with gr.Accordion("🐍 Python 範例", open=False):
                gr.Code(f'''import requests
import base64
from PIL import Image
import io

# API 設定
API_URL = "{SPACE_URL or "https://your-space.hf.space"}/api/generate"
API_KEY = "{PUBLIC_API_KEY}"

# 呼叫 API
response = requests.post(API_URL, json={{
    "prompt": "A beautiful sunset over ocean",
    "model": "black-forest-labs/FLUX.1-schnell",
    "api_key": API_KEY
}})

data = response.json()

if data["status"] == 200:
    # 解碼圖片
    img_data = base64.b64decode(data["image_base64"].split(",")[1])
    image = Image.open(io.BytesIO(img_data))
    image.save("output.png")
    print(f"✅ 圖片已保存！生成耗時: {{data['generation_time']}}s")
else:
    print(f"❌ 錯誤: {{data.get('error')}}")''', language="python")
            
            with gr.Accordion("🌐 JavaScript 範例", open=False):
                gr.Code(f'''// 呼叫 API 生成圖片
async function generateImage(prompt) {{
    const response = await fetch("{SPACE_URL or "https://your-space.hf.space"}/api/generate", {{
        method: 'POST',
        headers: {{ 'Content-Type': 'application/json' }},
        body: JSON.stringify({{
            prompt: prompt,
            model: "black-forest-labs/FLUX.1-schnell",
            api_key: "{PUBLIC_API_KEY}"
        }})
    }});
    
    const data = await response.json();
    
    if (data.status === 200) {{
        // 顯示圖片
        const img = document.createElement('img');
        img.src = data.image_base64;
        document.body.appendChild(img);
        return data;
    }} else {{
        console.error('Error:', data.error);
    }}
}}

// 使用
generateImage("A cute cat in space suit");''', language="javascript")
        
        # === Tab 3: 設定說明 ===
        with gr.Tab("⚙️ 設定說明"):
            gr.Markdown(f"""
            ### 環境變數設定
            
            請在 Space **Settings > Variables and Secrets** 設定：
            
            | 變數名稱 | 說明 | 必需 |
            |----------|------|------|
            | `HF_TOKEN` | Hugging Face Access Token | 否 |
            | `PUBLIC_API_KEY` | 公開 API Key（預設 `{PUBLIC_API_KEY}`）| 否 |
            | `SPACE_URL` | Space 完整 URL | 否 |
            
            ### 如何取得 HF_TOKEN
            
            1. 前往 [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)
            2. 建立 `Read` 權限的 Token
            3. 複製到 Space Settings > Secrets
            
            ### 使用限制
            
            - **免費帳號**：模型 < 10GB，Rate Limit 300 requests/hour
            - **Pro 帳號**：$10/月，1000 requests/hour
            
            ### 錯誤處理
            
            | HTTP 狀態碼 | 說明 |
            |-------------|------|
            | 200 | 成功 |
            | 400 | 請求參數錯誤 |
            | 401 | API Key 無效 |
            | 403 | 超過 10GB 限制或無權限 |
            | 500 | 伺服器內部錯誤 |
            """)

# 掛載 Gradio 到 FastAPI
from fastapi.middleware.cors import CORSMiddleware

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 使用 mount_gradio_app 將 Gradio 掛載到 FastAPI
import gradio as gr
app = gr.mount_gradio_app(app, demo, path="/")

# 啟動
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)