src/services/tryon.py

"""虛擬試穿服務 — Multi-provider 架構（透過 ModelRegistry）

【這個模組在做什麼？】
  接收使用者照片 + 衣服照片，產出試穿合成圖。
  Provider 設定從 ModelRegistry 的 models.yaml 讀取，由 TRYON_PROVIDER 環境變數切換。

【Provider 說明】
  leffa（預設）：   Meta Leffa via HF Space franciszzj/Leffa，MIT 商用安全
  huggingface：    HuggingFace Inference API 跑 IDM-VTON（免費層，受 CC BY-NC-SA 約束）
  replicate：      Replicate IDM-VTON API，$0.023/次
  fashn：          Fashn.ai API，商用安全（無 CC BY-NC-SA 限制）
  local：          本地 diffusers pipeline（需 GPU，學習/研究用）
  mock：           開發用假 URL

【授權警告】
  IDM-VTON 是 CC BY-NC-SA 4.0（非商用）—— huggingface / replicate / local 三條路都受約束。
  Leffa 是 MIT（商用安全），預設選用。
  Fashn.ai 另有商用授權，作為商業級替代。

【ModelRegistry 整合】
  tryon capability → idm-vton model（provider: external）—— model 名稱是歷史命名，
  實際上 variants 裡混合多個 VTON 模型（IDM-VTON / Leffa / Fashn）。
  active_variant 由 TRYON_PROVIDER 環境變數決定（預設 leffa）。
"""

from __future__ import annotations

import asyncio
import logging
import os

import httpx
from pydantic import BaseModel

logger = logging.getLogger(__name__)


class TryOnResult(BaseModel):
    result_image_url: str
    provider: str = "replicate_idm_vton"


# ============================================================================
# ModelRegistry 整合：讀取 active_variant 和 variant config
# ============================================================================


async def _get_tryon_config() -> tuple[str, dict]:
    """
    從 ModelRegistry 取得試穿設定。

    【回傳值】
    (active_variant, variant_config)
    - active_variant：當前選用的 provider 名稱（huggingface/replicate/fashn/local/mock）
    - variant_config：該 provider 的設定 dict（來自 models.yaml idm-vton.variants）

    【Fallback 邏輯】
    若 ModelRegistry 未初始化（例如單元測試），直接讀取環境變數。
    """
    try:
        from src.registry import get_registry

        reg = get_registry()
        # 讀取 capability 設定
        cap_config = reg.get_capability_config("tryon")
        # active_variant 由 cap_config["active_variant"] 取得，
        # 這個值在 models.yaml 中定義為 "${TRYON_PROVIDER:-leffa}"。
        # Fallback 一律用 leffa（商用安全）——IDM-VTON 的 huggingface 變體是 CC BY-NC-SA，
        # 任何「schema drift 或 registry race 時偷偷走回舊路徑」的縫隙都違反 F1 意圖。
        active_variant = cap_config.get("active_variant", "leffa")
        # 讀取 idm-vton model 設定（含 variants）
        model_config = reg.get_model_config(cap_config["model"])
        variant_config = model_config.get("variants", {}).get(active_variant, {})
        return active_variant, variant_config
    except Exception:
        # Fallback：registry 未初始化或 raise 時——仍預設 leffa，避免跳到 huggingface
        # （= IDM-VTON CC BY-NC-SA）違反 F1 的「commercial-safe by default」意圖。
        provider = os.environ.get("TRYON_PROVIDER", "leffa")
        return provider, {}


async def run_tryon(user_photo_url: str, garment_image_url: str) -> TryOnResult:
    """
    虛擬試穿入口。根據 ModelRegistry 設定的 active_variant 分發到對應 provider。

    【Provider 選擇邏輯】
    leffa → HF Space franciszzj/Leffa（MIT，預設，商用安全）
    huggingface → HuggingFace Inference API（需 HF_TOKEN，免費）
    replicate → Replicate API（需 REPLICATE_API_TOKEN，付費）
    fashn → Fashn.ai API（需 FASHN_API_KEY，付費，商用安全）
    local → 本地 diffusers（需 GPU + 模型下載）
    mock → 開發用假 URL
    其他/無 key → mock
    """
    # 從 registry 讀取 active_variant（而非直接讀環境變數）
    active_variant, variant_config = await _get_tryon_config()

    if active_variant == "leffa":
        return await _leffa_tryon(user_photo_url, garment_image_url, variant_config)
    elif active_variant == "huggingface":
        return await _huggingface_tryon(
            user_photo_url, garment_image_url, variant_config
        )
    elif active_variant == "fashn":
        return await _fashn_tryon(user_photo_url, garment_image_url, variant_config)
    elif active_variant == "local":
        return await _local_tryon(user_photo_url, garment_image_url)
    elif active_variant == "mock":
        return _mock_tryon(user_photo_url, garment_image_url)
    else:  # replicate (fallback)
        return await _replicate_tryon(user_photo_url, garment_image_url, variant_config)


# ============================================================================
# Leffa Provider（HF Space franciszzj/Leffa — MIT 商用安全，預設）
# ============================================================================


async def _leffa_tryon(
    user_photo_url: str,
    garment_image_url: str,
    variant_config: dict,
) -> TryOnResult:
    """
    Leffa — 透過 HuggingFace Space `franciszzj/Leffa` 跑 SOTA 虛擬試穿。

    【為什麼用 Leffa？】
    Meta 2024 年的模型（CVPR 2025，arXiv:2412.08486），MIT 授權——解開
    IDM-VTON CC BY-NC-SA 商用禁令。VITON-HD 成績贏過 IDM-VTON / OOTDiffusion，
    基於 SD-1.5 UNet 較輕（相對於 IDM-VTON 雙 pass），A100 約 6s/image；
    走 HF Space 共享 GPU 時 ~20-40s，冷啟動可能要多等 30-60s。

    【執行位置】
    呼叫 Space 託管的 Gradio endpoint—— 完全不用本地 GPU。Space 閒置會睡，
    第一次呼叫會等 ZeroGPU 喚醒。

    【variant_config（來自 models.yaml idm-vton.variants.leffa）】
      space_id:      "franciszzj/Leffa"
      api_name:      "/leffa_predict_vt"（另一個 api_name 是 /leffa_predict_pt 做姿態轉換）
      model_type:    "viton_hd" | "dress_code"
      garment_type:  "upper_body" | "lower_body" | "dresses"
      step:          30（30-100，越高越細節但越慢）
      scale:         2.5（CFG，0.1-5.0）
      seed:          固定值可重現，或 -1 隨機

    【輸出】
    Space 回傳 tuple (generated_image, mask, densepose)，我們只取第一張。
    gradio_client 把檔案下載到本地 temp 路徑，需讀成 base64 data URI 回前端。

    【async 規範】
    gradio_client 內部是 blocking HTTP + SSE polling，用 asyncio.to_thread 包起來
    避免凍結 event loop（遵守 ml-service async discipline rule）。
    """
    # SSRF guard: `handle_file(url)` 會在 gradio_client 內部對 URL 發 HTTP GET
    # 把檔案拉下來存成 local temp path，這筆下載發生在 ml-service 進程內。所以
    # 任何未經驗證的 userPhotoUrl / garmentImageUrl 都是一條 SSRF vector，可達
    # 169.254.169.254 / 127.0.0.1 / Redis 等。hub-service 的 /tryon route
    # 目前只跑 z.string().url()（不是 isSafeContentUrl），因此這道檢查必須
    # 在呼叫 handle_file 之前做。任一 URL 被擋 = RuntimeError → worker 寫
    # tryonJobs.status="failed" + 通知使用者。
    from src.ssrf import validate_image_url

    try:
        # offload: validate_image_url runs blocking socket.getaddrinfo (ML-03);
        # called from async _leffa_tryon, must to_thread or event loop stalls.
        await asyncio.to_thread(validate_image_url, user_photo_url)
        await asyncio.to_thread(validate_image_url, garment_image_url)
    except ValueError as e:
        raise RuntimeError(f"Leffa refused unsafe image URL: {e}") from e

    # 放 closure 內延遲 import—— 避免 ml-service 啟動時因 gradio_client 初始化成本
    # 影響 /health 回應（ml-service async discipline rule 的 C-extension pattern）
    from gradio_client import Client, handle_file

    space_id = variant_config.get("space_id", "franciszzj/Leffa")
    api_name = variant_config.get("api_name", "/leffa_predict_vt")
    model_type = variant_config.get("model_type", "viton_hd")
    garment_type = variant_config.get("garment_type", "upper_body")
    step = int(variant_config.get("step", 30))
    scale = float(variant_config.get("scale", 2.5))
    seed = int(variant_config.get("seed", 42))

    # HF token：Space 公開可匿名呼叫，但帶 token 可用 Pro 額度避免排隊
    hf_token = os.environ.get("HUGGINGFACE_TOKEN", "") or os.environ.get("HF_TOKEN", "")

    def _blocking() -> str:
        """Gradio client 同步呼叫——包進 to_thread 跑。"""
        client_kwargs: dict = {}
        if hf_token:
            client_kwargs["hf_token"] = hf_token
        client = Client(space_id, **client_kwargs)

        # handle_file(url) 把 URL 下載到 temp 檔，符合 Gradio file component 輸入格式。
        # 注意：參數名是 `src_image_path` / `ref_image_path`（Space 的實際 kwarg，不是 `src_image`）。
        # ref_acceleration / vt_repaint 是 Radio[True, False]，view_api 顯示的 Literal['True','False']
        # 是 **label** 文字，實際 choices 收的是 Python bool。
        result = client.predict(
            src_image_path=handle_file(user_photo_url),
            ref_image_path=handle_file(garment_image_url),
            ref_acceleration=False,
            step=step,
            scale=scale,
            seed=seed,
            vt_model_type=model_type,
            vt_garment_type=garment_type,
            vt_repaint=False,
            api_name=api_name,
        )
        # Leffa 回 tuple (generated_image, mask, densepose)，取第一張
        first = result[0] if isinstance(result, (list, tuple)) else result
        # gradio_client 可能回 local temp path string 或 {"path": ..., "url": ...} dict
        if isinstance(first, dict):
            return str(first.get("url") or first.get("path") or "")
        return str(first)

    logger.info(
        "Leffa via HF Space %s (model=%s, garment=%s, step=%d, scale=%.1f)",
        space_id,
        model_type,
        garment_type,
        step,
        scale,
    )
    try:
        result_path = await asyncio.to_thread(_blocking)
    except Exception as e:
        logger.error("Leffa HF Space call failed: %s", e)
        raise RuntimeError(
            f"Leffa HF Space request failed ({type(e).__name__}): {e}. "
            "Check space availability at https://huggingface.co/spaces/franciszzj/Leffa "
            "or switch TRYON_PROVIDER to fashn/mock."
        ) from e

    if not result_path:
        raise RuntimeError("Leffa returned empty result path")

    # 遠端 URL 直接回傳；local temp file 轉 data URI 讓前端直接顯示
    if result_path.startswith(("http://", "https://")):
        result_url = result_path
    else:
        import base64
        import mimetypes

        with open(result_path, "rb") as f:
            img_bytes = f.read()
        # Gradio 會存成 .webp / .png 等，依副檔名決定 mime；未知則回 webp（Leffa 現況）
        mime, _ = mimetypes.guess_type(result_path)
        mime = mime or "image/webp"
        b64 = base64.b64encode(img_bytes).decode("utf-8")
        result_url = f"data:{mime};base64,{b64}"

    return TryOnResult(result_image_url=result_url, provider="leffa_hf_space")


# ============================================================================
# HuggingFace Inference API Provider（免費，CC BY-NC-SA 非商用）
# ============================================================================


async def _huggingface_tryon(
    user_photo_url: str,
    garment_image_url: str,
    variant_config: dict,
) -> TryOnResult:
    """
    HuggingFace Inference API — 免費層虛擬試穿。

    【學習重點：HuggingFace Inference API 的運作方式】
    HuggingFace 提供免費的 Inference API，讓你不用自己架 GPU 就能呼叫模型。
    原理是 HuggingFace 在雲端幫你跑模型推論（inference），你只要 HTTP POST 就好。

    【API 流程】
    1. 取得 HF_TOKEN（在 https://huggingface.co/settings/tokens 建立，免費）
    2. POST JSON 到 https://api-inference.huggingface.co/models/{model_id}
    3. Header 帶 Authorization: Bearer {token}
    4. Body 放模型需要的輸入參數（每個模型不同）

    【免費層限制】
    - 模型可能需要「冷啟動」（cold start）：第一次呼叫要等模型載入（~30-60s）
    - API 回 503 + "estimated_time" 表示模型正在載入，需要重試
    - 免費層有 rate limit（約 30 req/min），超過會被暫時限流
    - 大模型（如 IDM-VTON）可能需要 Pro 帳號才能穩定使用

    【授權注意】
    IDM-VTON 是 CC BY-NC-SA 4.0，透過 HF Inference API 呼叫同樣受此約束。
    商用請改用 fashn provider。

    【ModelRegistry 參數】
    variant_config 來自 models.yaml idm-vton.variants.huggingface：
    - model_id: "yisol/IDM-VTON"
    """
    # 嘗試讀取 HuggingFace token（支援兩種環境變數名稱）
    # HUGGINGFACE_TOKEN 是完整名稱，HF_TOKEN 是 HuggingFace CLI 的慣例
    token = os.environ.get("HUGGINGFACE_TOKEN", "") or os.environ.get("HF_TOKEN", "")

    if not token:
        # 沒有 token 就 fallback 到 mock（不要讓服務掛掉）
        logger.warning(
            "HF_TOKEN / HUGGINGFACE_TOKEN not set, falling back to mock. "
            "Get a free token at https://huggingface.co/settings/tokens"
        )
        return _mock_tryon(user_photo_url, garment_image_url)

    # 從 variant_config 取得 model_id，fallback 到預設值
    model_id = variant_config.get("model_id", "yisol/IDM-VTON")

    # HuggingFace Inference API 的端點格式：固定前綴 + 模型 ID
    api_url = f"https://api-inference.huggingface.co/models/{model_id}"
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
    }

    # 【學習重點：模型輸入格式】
    # 不同模型接受的 JSON 結構不同，要看模型的 API 文件。
    # IDM-VTON 需要人物照片（human_img）和衣服照片（garm_img）。
    payload = {
        "inputs": {
            "human_img": user_photo_url,
            "garm_img": garment_image_url,
        },
    }

    # 最多重試 3 次（處理模型冷啟動的 503 回應）
    max_retries = 3

    async with httpx.AsyncClient(timeout=180.0) as client:
        for attempt in range(max_retries):
            logger.info(
                "HuggingFace Inference API attempt %d/%d for model %s",
                attempt + 1,
                max_retries,
                model_id,
            )

            resp = await client.post(api_url, headers=headers, json=payload)

            # 【學習重點：503 = 模型正在載入（冷啟動）】
            # HF 免費層的模型不會一直跑在 GPU 上，閒置一段時間後會被卸載。
            # 第一次呼叫時 HF 要重新載入模型到 GPU，這段時間 API 回 503。
            # 回應的 JSON 裡有 "estimated_time" 告訴你大概要等多久。
            if resp.status_code == 503:
                body = resp.json()
                wait_time = body.get("estimated_time", 30.0)
                # 上限 60 秒，避免等太久
                wait_time = min(float(wait_time), 60.0)
                logger.info(
                    "Model is loading (cold start), waiting %.1fs before retry...",
                    wait_time,
                )
                await asyncio.sleep(wait_time)
                continue

            # 【學習重點：429 = 超過 rate limit】
            # 免費層有請求頻率限制，超過時 API 回 429。
            if resp.status_code == 429:
                logger.warning(
                    "HuggingFace rate limit hit, waiting 10s before retry..."
                )
                await asyncio.sleep(10)
                continue

            # 其他錯誤直接拋出
            resp.raise_for_status()

            # 【學習重點：回應格式】
            # HF Inference API 對圖片模型可能回傳：
            #   (a) JSON：包含 base64 編碼的圖片或 URL
            #   (b) 二進位：直接回傳圖片 bytes（Content-Type: image/...）
            # 我們需要根據 Content-Type 判斷回應類型。
            content_type = resp.headers.get("content-type", "")

            if "application/json" in content_type:
                # JSON 回應：嘗試取出圖片 URL 或 base64
                data = resp.json()
                if isinstance(data, list) and len(data) > 0:
                    # 有些模型回傳陣列格式
                    result_url = str(data[0])
                elif isinstance(data, dict):
                    # 嘗試常見的回應欄位名稱
                    result_url = data.get(
                        "output", data.get("image", data.get("url", ""))
                    )
                    if not result_url:
                        # 如果有 base64 資料，轉成 data URI
                        b64 = data.get("generated_image", data.get("image_base64", ""))
                        if b64:
                            result_url = f"data:image/png;base64,{b64}"
                        else:
                            raise RuntimeError(
                                f"HuggingFace returned unexpected JSON structure: {list(data.keys())}"
                            )
                else:
                    result_url = str(data)

                return TryOnResult(
                    result_image_url=result_url,
                    provider="huggingface_idm_vton",
                )

            elif content_type.startswith("image/"):
                # 二進位圖片回應：轉成 base64 data URI
                # 【學習重點：Data URI】
                # 格式是 data:{mime_type};base64,{base64_data}
                # 瀏覽器可以直接顯示，不需要額外的圖片 URL
                import base64

                img_b64 = base64.b64encode(resp.content).decode("utf-8")
                data_uri = f"data:{content_type};base64,{img_b64}"
                return TryOnResult(
                    result_image_url=data_uri,
                    provider="huggingface_idm_vton",
                )

            else:
                raise RuntimeError(
                    f"HuggingFace returned unexpected content-type: {content_type}"
                )

        # 所有重試都失敗（通常是模型一直載入不完）
        raise RuntimeError(
            "HuggingFace Inference API failed after all retries. "
            "Model may be too large for free tier — consider upgrading to HF Pro "
            "or switching to TRYON_PROVIDER=replicate."
        )


# ============================================================================
# Replicate Provider
# ============================================================================


async def _replicate_tryon(
    user_photo_url: str,
    garment_image_url: str,
    variant_config: dict,
) -> TryOnResult:
    """Replicate IDM-VTON — 直接 HTTP 呼叫（不用 replicate package）。

    【學習重點：非同步 API 呼叫模式】
    1. POST 建立 prediction（拿到 prediction ID）
    2. 輪詢（polling）GET prediction 狀態，等 succeeded/failed
    3. 用 asyncio.sleep 不阻塞 event loop（不用 time.sleep！）
    4. httpx.AsyncClient 是 Python 的非同步 HTTP client（類似 Node 的 fetch）

    【ModelRegistry 參數】
    variant_config 來自 models.yaml idm-vton.variants.replicate：
    - model_id: "cuuupid/idm-vton"
    """
    api_token = os.environ.get("REPLICATE_API_TOKEN", "")
    is_real_key = len(api_token) > 20 and not api_token.startswith("your")

    if not is_real_key:
        return _mock_tryon(user_photo_url, garment_image_url)

    # 註：replicate 用 version hash 鎖模型版本（payload["version"]），model_id
    # 欄位純 documentation 用途，因此 variant_config["model_id"] 不在此讀取。
    headers = {
        "Authorization": f"Bearer {api_token}",
        "Content-Type": "application/json",
    }
    payload = {
        "version": "0513734a452173b8173e907e3a59d19a36266e55b48528559432bd21c7d7e985",
        "input": {
            "human_img": user_photo_url,
            "garm_img": garment_image_url,
        },
    }

    async with httpx.AsyncClient(timeout=120.0) as client:
        resp = await client.post(
            "https://api.replicate.com/v1/predictions",
            headers=headers,
            json=payload,
        )
        if resp.status_code == 402:
            raise RuntimeError(
                "Replicate billing required: add a payment method at https://replicate.com/account/billing"
            )
        resp.raise_for_status()
        prediction = resp.json()

        poll_url = prediction["urls"]["get"]
        for _ in range(120):
            poll_resp = await client.get(poll_url, headers=headers)
            poll_resp.raise_for_status()
            data = poll_resp.json()
            status = data["status"]

            if status == "succeeded":
                output = data["output"]
                result_url = str(output[0]) if isinstance(output, list) else str(output)
                return TryOnResult(
                    result_image_url=result_url, provider="replicate_idm_vton"
                )
            elif status in ("failed", "canceled"):
                error_msg = data.get("error", "Unknown error")
                raise RuntimeError(f"Replicate prediction {status}: {error_msg}")

            await asyncio.sleep(1)

        raise RuntimeError("Replicate prediction timed out after 120s")


# ============================================================================
# Fashn.ai Provider（商用安全替代）
# ============================================================================


async def _fashn_tryon(
    user_photo_url: str,
    garment_image_url: str,
    variant_config: dict,
) -> TryOnResult:
    """
    Fashn.ai — 商用安全的虛擬試穿 API。

    【為什麼用 Fashn.ai？】
    IDM-VTON 是 CC BY-NC-SA 4.0（非商用），商業化前必須換。
    Fashn.ai 提供商用授權的試穿 API，品質接近 IDM-VTON。

    API 文件：https://docs.fashn.ai

    【ModelRegistry 參數】
    variant_config 來自 models.yaml idm-vton.variants.fashn：
    - api_url: "https://api.fashn.ai/v1"
    """
    api_key = os.environ.get("FASHN_API_KEY", "")
    if not api_key:
        logger.warning("FASHN_API_KEY not set, falling back to mock")
        return _mock_tryon(user_photo_url, garment_image_url)

    # 從 variant_config 讀取 API URL，fallback 到預設
    api_url = variant_config.get("api_url", "https://api.fashn.ai/v1")

    async with httpx.AsyncClient(timeout=120.0) as client:
        # Fashn.ai run endpoint
        resp = await client.post(
            f"{api_url}/run",
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            json={
                "model_image": user_photo_url,
                "garment_image": garment_image_url,
                "category": "auto",
            },
        )
        resp.raise_for_status()
        data = resp.json()

        # Poll for result
        pred_id = data.get("id", "")
        for _ in range(120):
            status_resp = await client.get(
                f"{api_url}/status/{pred_id}",
                headers={"Authorization": f"Bearer {api_key}"},
            )
            status_resp.raise_for_status()
            status_data = status_resp.json()

            if status_data.get("status") == "completed":
                return TryOnResult(
                    result_image_url=status_data["output"],
                    provider="fashn_ai",
                )
            elif status_data.get("status") in ("failed", "canceled"):
                raise RuntimeError(
                    f"Fashn.ai prediction failed: {status_data.get('error')}"
                )

            await asyncio.sleep(1)

        raise RuntimeError("Fashn.ai prediction timed out")


# ============================================================================
# Local Provider（學習/研究用，需 GPU）
# ============================================================================


async def _local_tryon(user_photo_url: str, garment_image_url: str) -> TryOnResult:
    """
    本地 IDM-VTON — diffusers pipeline（需 NVIDIA GPU 16GB+）。

    【學習重點：HuggingFace diffusers 的用法】
    真正實作時的程式碼大致是：
    ```python
    from diffusers import AutoPipelineForImage2Image
    pipe = AutoPipelineForImage2Image.from_pretrained("yisol/IDM-VTON")
    pipe = pipe.to("cuda")  # 需要 NVIDIA GPU
    result = pipe(person_image, garment_image)
    result.images[0].save("output.png")
    ```

    【注意】
    - CC BY-NC-SA 4.0：僅限學習/研究，禁止商用
    - 需要預先下載模型（~10GB）
    - 第一次推論較慢（模型載入 ~30s）

    【誠實不可用 — feedback_no_fabricated_fallbacks】
    本地 diffusers pipeline 從未實作（需租 GPU 實測後才填入真正的程式碼）。
    在那之前，這個 provider 不可用——以前回傳假 mock URL 是捏造成功，會讓上層
    以為試穿真的生成了。改成 raise RuntimeError（與 leffa/replicate/fashn 失敗時
    一致的誠實路徑）：上層 worker 收到後寫 tryonJobs.status="failed" + 通知使用者，
    HTTP 路徑回 500/503，前端走既有的 honest-unavailable 顯示而非假成功。
    商用請用 leffa（預設）或 fashn provider。
    """
    raise RuntimeError(
        "Local IDM-VTON provider is not implemented (requires a GPU). "
        "Use the leffa (default) or fashn provider instead."
    )


# ============================================================================
# Mock（開發用）
# ============================================================================


def _mock_tryon(user_photo_url: str, garment_image_url: str) -> TryOnResult:
    """開發用 mock：回傳假 URL。"""
    mock_id = abs(hash(user_photo_url + garment_image_url)) % 100000
    mock_url = f"https://mock-tryon.wardrobe-os.dev/results/{mock_id}.jpg"
    return TryOnResult(result_image_url=mock_url, provider="mock")