#!/usr/bin/env python3
"""
InsightFace 证件照换脸 API 服务(全功能版)
启动: python3 idphoto_server.py [--port 7870]

端点:
  /health                健康检查
  /swap       POST bash  换头(保留头发, 仿射+inpaint)  - 备选
  /swapface   POST bash  精确换脸(inswapper, 只换五官)
  /idphoto    POST json  证件照合成(单人像->套模板->换脸)
"""
import os
import sys
import glob
import base64
import argparse
import time
import threading
from io import BytesIO
from contextlib import asynccontextmanager

import cv2
import numpy as np
from PIL import Image
from fastapi import FastAPI
from pydantic import BaseModel

BASE = os.path.dirname(os.path.abspath(__file__))
TPL_DIR = os.path.join(BASE, "templates")
INSIGHTFACE_ROOT = os.path.join(BASE, "models", "insightface")
INSIGHTFACE_MODELS = os.path.join(INSIGHTFACE_ROOT, "models")

app_analyzer = None
app_swapper = None
load_error = None
_templates = {}


def _load():
    global app_analyzer, app_swapper, load_error
    try:
        print("[load] initializing InsightFace...", flush=True)
        from insightface.app import FaceAnalysis
        app_analyzer = FaceAnalysis(
            name="buffalo_l",
            root=INSIGHTFACE_ROOT,
            providers=["CPUExecutionProvider"],
            allowed_modules=["detection", "landmark_2d_106", "genderage",
                             "recognition"],
        )
        app_analyzer.prepare(ctx_id=0, det_size=(640, 640))

        import insightface
        swapper_path = os.path.join(BASE, "inswapper_128.onnx")
        if not os.path.exists(swapper_path):
            swapper_path = os.path.join(
                INSIGHTFACE_MODELS, "inswapper_128.onnx")
        app_swapper = insightface.model_zoo.get_model(
            swapper_path, providers=["CPUExecutionProvider"])
        _scan_templates()
        print("[load] READY", flush=True)
    except Exception as e:
        import traceback
        traceback.print_exc()
        load_error = f"{type(e).__name__}: {e}"
        print(f"[load] FAILED: {load_error}", flush=True)


def _scan_templates():
    _templates.clear()
    for f in glob.glob(os.path.join(TPL_DIR, "*", "*.jpg")) + \
            glob.glob(os.path.join(TPL_DIR, "*", "*.png")):
        rel = os.path.relpath(f, TPL_DIR)
        gender_dir = os.path.dirname(rel)
        name = os.path.splitext(os.path.basename(f))[0]
        _templates.setdefault(gender_dir, {})[name] = f
    print(f"[load] templates: {_templates}", flush=True)


@asynccontextmanager
async def lifespan(_app):
    t = threading.Thread(target=_load, daemon=True)
    t.start()
    yield


api = FastAPI(title="InsightFace ID Photo", lifespan=lifespan)


class SwapRequest(BaseModel):
    src_image: str
    dst_image: str
    expand: float = 2.0
    quality: int = 95


class IDPhotoRequest(BaseModel):
    face_image: str          # 上传的人像 base64
    gender: str = "auto"     # auto|male|female
    color: str = "blue"      # white|blue|red (证件照底色)
    quality: int = 95


def _decode_b64(b64_str: str) -> np.ndarray:
    data = base64.b64decode(b64_str)
    img = Image.open(BytesIO(data)).convert("RGB")
    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)


def _encode_b64(img: np.ndarray, quality: int) -> str:
    ok, enc = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, quality])
    if not ok:
        raise ValueError("encode failed")
    return base64.b64encode(enc.tobytes()).decode()


def _pick_face(faces, prefer=0):
    if not faces:
        return None
    if prefer < len(faces):
        return faces[prefer]
    return faces[0]


def _swap_head(src, dst, expand):
    """仿射+inpaint, 整个头(含发)替换"""
    src_faces = app_analyzer.get(src)
    dst_faces = app_analyzer.get(dst)
    if not src_faces:
        raise ValueError("source: no face")
    if not dst_faces:
        raise ValueError("target: no face")
    src_kps = src_faces[0].kps
    dst_kps = dst_faces[0].kps

    s = src_kps.astype(np.float32)
    d = dst_kps.astype(np.float32)
    M, _ = cv2.estimateAffinePartial2D(s, d, method=cv2.LMEDS)
    if M is None:
        raise ValueError("affine failed")

    # 掩膜(椭圆包住头+发)
    cx = src_kps[:, 0].mean(); cy = src_kps[:, 1].mean()
    fw = src_kps[:, 0].max() - src_kps[:, 0].min()
    fh = src_kps[:, 1].max() - src_kps[:, 1].min()
    rx = fw * expand / 2
    ry = fh * expand * 1.1 / 2
    h, w = dst.shape[:2]
    mask = np.zeros((h, w), dtype=np.uint8)
    cv2.ellipse(mask, (int(cx), int(cy - ry * 0.08)),
                (int(rx), int(ry)), 0, 360, 255, -1)
    blur_k = max(31, int(min(rx, ry) * 0.4) | 1)
    mask = cv2.GaussianBlur(mask, (blur_k, blur_k), blur_k // 3)

    src_warped = cv2.warpAffine(src, M, (w, h), flags=cv2.INTER_LINEAR,
                                borderMode=cv2.BORDER_CONSTANT, borderValue=0)
    mask_warped = cv2.warpAffine(mask, M, (w, h), flags=cv2.INTER_LINEAR,
                                 borderMode=cv2.BORDER_CONSTANT, borderValue=0)

    erase = (mask_warped > 20).astype(np.uint8) * 255
    erase = cv2.dilate(erase, np.ones((7, 7), np.uint8), iterations=2)
    dst_erased = cv2.inpaint(dst, (erase > 128).astype(np.uint8) * 255,
                             inpaintRadius=10, flags=cv2.INPAINT_TELEA)

    alpha = (mask_warped.astype(np.float32) / 255.0)[:, :, np.newaxis]

    # 色调匹配
    outer = cv2.dilate(erase, np.ones((20, 20), np.uint8)) - erase
    shift = np.zeros(3, dtype=np.float32)
    if outer.sum() > 0 and (mask_warped > 128).sum() > 0:
        op = dst[outer > 0].astype(np.float32).mean(axis=0)
        ip = src_warped[mask_warped > 128].astype(np.float32).mean(axis=0)
        shift = (op - ip) * 0.35
    src_adj = np.clip(src_warped.astype(np.float32) + shift, 0, 255)
    result = dst_erased.astype(np.float32) * (1 - alpha) + src_adj * alpha
    return np.clip(result, 0, 255).astype(np.uint8)


def _swap_face(dst, src_face, dst_face):
    """inswapper 精确换脸(只换五官, 自然融合)"""
    return app_swapper.get(dst, dst_face, src_face, paste_back=True)


def _detect(gender, img, name):
    faces = app_analyzer.get(img)
    if not faces:
        raise ValueError(f"{name}: no face detected")
    if gender == "auto":
        # 取最大人脸
        return max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
    return faces[0]


def _apply_bg_color(img, face, color):
    """把模板背景色统一替换(基于人脸外区域主色判定)"""
    COLORS = {"white": (255, 255, 255), "blue": (67, 142, 219), "red": (200, 40, 55),
              "blue2": (91, 155, 213), "red2": (190, 30, 45)}
    if color not in COLORS:
        return img
    h, w = img.shape[:2]
    # 简单处理: 四周背景采样中值色, 属于偏像素则替换(用阈值)
    border = np.concatenate([
        img[0:5, :, :].reshape(-1, 3),
        img[-5:, :, :].reshape(-1, 3),
        img[:, 0:5, :].reshape(-1, 3),
        img[:, -5:, :].reshape(-1, 3),
    ])
    bg = np.median(border, axis=0).astype(np.float32)
    # 基于人脸半径确定安全区, 以外区域若接近背景则涂色
    cx = (face.bbox[0] + face.bbox[2]) / 2
    cy = (face.bbox[1] + face.bbox[3]) / 2
    r = (face.bbox[2] - face.bbox[0]) * 3.2
    yy, xx = np.mgrid[0:h, 0:w]
    dist = np.sqrt((xx - cx) ** 2 + (yy - cy) ** 2)
    outside = dist > r
    diff = np.abs(img.astype(np.float32) - bg).mean(axis=2)
    bg_pixel = outside & (diff < 40)
    target = np.array(COLORS[color], dtype=np.float32)
    out = img.astype(np.float32).copy()
    mix = np.zeros((h, w), dtype=np.float32)
    mix[bg_pixel] = 1.0
    # 边缘柔化
    mix = cv2.GaussianBlur(mix, (51, 51), 20)
    mix3 = mix[:, :, np.newaxis]
    out = out * (1 - mix3) + target * mix3
    return np.clip(out, 0, 255).astype(np.uint8)


@api.get("/health")
def health():
    if app_analyzer is not None:
        return {"status": "ok", "model": "buffalo_l", "engine": "insightface"}
    if load_error:
        return {"status": "error", "error": load_error}
    return {"status": "loading"}


@api.post("/swap")
def swap(req: SwapRequest):
    if app_analyzer is None:
        return {"error": "not loaded"}
    t0 = time.time()
    try:
        src = _decode_b64(req.src_image)
        dst = _decode_b64(req.dst_image)
        result = _swap_head(src, dst, req.expand)
        b64 = _encode_b64(result, req.quality)
        return {"image": b64, "time_ms": int((time.time() - t0) * 1000)}
    except Exception as e:
        return {"error": f"{type(e).__name__}: {e}"}


@api.post("/swapface")
def swapface(req: SwapRequest):
    """inswapper 精确换脸: src 的人脸 -> dst 的人脸位置"""
    if app_analyzer is None or app_swapper is None:
        return {"error": "not loaded"}
    t0 = time.time()
    try:
        src = _decode_b64(req.src_image)
        dst = _decode_b64(req.dst_image)
        src_fc = app_analyzer.get(src)
        dst_fc = app_analyzer.get(dst)
        if not src_fc or not dst_fc:
            raise ValueError("face not detected")
        result = _swap_face(dst, src_fc[0], dst_fc[0])
        b64 = _encode_b64(result, req.quality)
        return {"image": b64, "time_ms": int((time.time() - t0) * 1000)}
    except Exception as e:
        return {"error": f"{type(e).__name__}: {e}"}


@api.post("/idphoto")
def idphoto(req: IDPhotoRequest):
    """证件照合成: 单人像 -> swapper换到模板 -> 可选改底色"""
    if app_analyzer is None or app_swapper is None:
        return {"error": "not loaded"}
    t0 = time.time()
    try:
        face = _decode_b64(req.face_image)
        # 性别
        fc = app_analyzer.get(face)
        if not fc:
            return {"error": "no face in input"}
        fg = fc[0]
        guess = "male" if fg.sex == 1 else "female"
        gender = req.gender if req.gender in ("male", "female") else guess

        # 找模板
        gdir = _templates.get(gender)
        if not gdir:
            return {"error": f"no {gender} template"}
        tname = list(gdir.keys())[0]
        tpath = gdir[tname]
        tpl = cv2.imread(tpath)
        tfc = app_analyzer.get(tpl)
        if not tfc:
            return {"error": "template: no face"}

        result = _swap_face(tpl, fg, tfc[0])
        result = _apply_bg_color(result, app_analyzer.get(result)[0], req.color)
        b64 = _encode_b64(result, req.quality)
        return {"image": b64, "time_ms": int((time.time() - t0) * 1000),
                "gender": gender, "template": tpath, "input_gender": guess,
                "templates": list(_templates.keys())}
    except Exception as e:
        return {"error": f"{type(e).__name__}: {e}"}


@api.get("/")
def root():
    return {
        "service": "insightface-idphoto",
        "endpoints": ["/health", "/swap", "/swapface", "/idphoto", "/docs"],
        "templates": _templates,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=7870)
    parser.add_argument("--host", default="127.0.0.1")
    args = parser.parse_args()
    import uvicorn
    uvicorn.run(api, host=args.host, port=args.port)


if __name__ == "__main__":
    main()