import os
import io
import base64
import traceback
from PIL import Image
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse

from huggingface_hub import hf_hub_download
from pipeline import IPFacePlusV2Pipeline

# ✅ Hugging Faceのキャッシュ・設定ディレクトリを/tmpに変更
os.environ["HF_HOME"] = "/tmp/hf"
os.environ["HF_HUB_CACHE"] = "/tmp/hf/hub"
os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp/hf/hub"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf/transformers"
os.environ["XDG_CACHE_HOME"] = "/tmp/hf/xdg"

# ✅ トークンを取得（SpacesのSecretsに設定してあるもの）
hf_token = os.environ.get("HUGGINGFACE_HUB_TOKEN")
if not hf_token:
    raise ValueError("環境変数 'HUGGINGFACE_HUB_TOKEN' が設定されていません。")

# ✅ モデルファイルのダウンロード
model_path = hf_hub_download(
    repo_id="revi13/ip-adapter-faceid-private",
    filename="ip-adapter-faceid-plusv2_sd15.bin",
    local_dir="/tmp/hf/models",  # 保存先を手動指定
    local_dir_use_symlinks=False,  # シンボリックリンクではなく実体保存
    token=os.environ["HUGGINGFACE_HUB_TOKEN"]
)

# ✅ パイプライン初期化
pipeline = IPFacePlusV2Pipeline(model_path)

# ✅ FastAPI アプリ作成
app = FastAPI()

@app.get("/")
def root():
    return {"status": "ok"}

@app.post("/generate")
async def generate_image_api(request: Request):
    print("✅ [START] /generate API called")  # ← 呼び出し確認
    
    try:
        body = await request.json()
        base64_image = body.get("image")
        prompt = body.get("prompt", "portrait")
        print(f"🔧 Received prompt: {prompt[:30]}...")  # プロンプト確認

        print("🚀 Start generating image...")
        
        image_base64 = pipeline.generate_image(base64_image, prompt)

        print("✅ Image generation complete")
        
        return JSONResponse(content={"success": True, "image_base64": image_base64})
    except Exception as e:
        print(f"❌ Exception occurred: {str(e)}")
        print("=== ERROR OCCURRED ===")
        traceback.print_exc()
        print("======================")
        return JSONResponse(status_code=500, content={"success": False, "error": str(e)})