"""dfm_swapper.py — 把 DeepFaceLive 的 .dfm（DeepFaceLab 训练出的**专属人物模型**）接进实时管线。

为什么要它（CLAUDE.md §11.26–§11.28）：一次性换脸（inswapper 一族）只转移五官纹理，脸型/眉形/胡子/下颌
永远是目标本人的；专属模型整张脸都是那个人，表情幅度也更完整（keanu_reeves_320 实测张嘴 0.72 vs inswapper 0.55）。

接入口径逐字对照 FaceFusion `processors/modules/deep_swapper/core.py`（master 2026-09-07），不是自己发明：
  · 裁剪模板 `dfl_whole_face`（5 点 → cv2.estimateAffinePartial2D RANSAC 阈 100）@ 模型分辨率（224/320/384/448）
  · 预处理 = 反锐化 addWeighted(1.75, GaussianBlur σ2, −0.75) 再 /255，NHWC float32
  · 输入 `in_face:0`；AMP 架构另有 `morph_value:0`（0~1，1=全量像模型人物）
  · 输出 `out_face_mask:0`（目标脸区）/ `out_celeb_face:0`（换好的脸）/ `out_celeb_face_mask:0`（生成脸区）
  · 遮罩 = 盒状(blur 0.3) ∧ min(两 mask) 腐蚀3 + 高斯 6.25；再 ∧ 遮挡感知（我们的 xseg）
与 RoiSwapper 保持同一调用契约：`sw(img, face) -> img`，属性 occluder / occ_every / _occ_cache / edge_converge，
所以 pipeline 只需按人设类型二选一，修复器照旧拿 `_occ_cache`（arcface-128 对齐空间）。

模型来源：内置清单 `dfm_models.json`（随代码），用户添加的记在 `assets/dfm_user.json`（数据，不随代码同步）。
模型文件放 `assets/<name>.dfm`（0.5–2GB，按需下载，不走 scp 同步），预览图 `assets/dfm_<name>.jpg`。
"""
import json
import os
import sys
import time
import threading
import urllib.request

import cv2
import numpy as np

from colorfit import local_color_transfer

HERE = os.path.dirname(os.path.abspath(__file__))
ASSETS = os.path.join(HERE, "assets")
REGISTRY = os.path.join(HERE, "dfm_models.json")
USER_REG = os.path.join(ASSETS, "dfm_user.json")
PREFIX = "dfm:"                                   # 人设名前缀：persona_name = "dfm:<model_name>"
# FaceFusion face_helper.WARP_TEMPLATES['dfl_whole_face']
TEMPLATE = np.array([[0.35342266, 0.39285716], [0.62797622, 0.39285716], [0.48660713, 0.54017860],
                     [0.38839287, 0.68750011], [0.59821427, 0.68750011]], np.float32)
_LOCK = threading.Lock()


def _log(*a):
    s = "[dfm] " + " ".join(str(x) for x in a)
    print(s, flush=True)
    try:                                          # 与 pipeline_rt._log 同一份日志文件（§10.12：Vast 日志看不到管线侧 print）
        with open(os.path.join(HERE, "pipeline.log"), "a", encoding="utf-8") as f:
            f.write(time.strftime("%H:%M:%S ") + s + "\n")
    except Exception:
        pass


# ---------------------------------------------------------------- 清单
def _read_json(path):
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None


def is_model_name(persona_name):
    return isinstance(persona_name, str) and persona_name.startswith(PREFIX)


def model_name_of(persona_name):
    return persona_name[len(PREFIX):] if is_model_name(persona_name) else None


def safe_name(name):
    """只允许 [a-z0-9_-]，其余替换成 _；防止路径穿越与非 ASCII 路径（§5 cv2 中文路径坑）。"""
    name = os.path.splitext(os.path.basename(str(name)))[0].lower()
    return "".join(c if (c.isalnum() or c in "_-") else "_" for c in name)[:64] or "model"


def model_path(name):
    return os.path.join(ASSETS, safe_name(name) + ".dfm")


def preview_file(name):
    return "dfm_" + safe_name(name) + ".jpg"


def _builtin():
    d = _read_json(REGISTRY) or {}
    return [m for m in d.get("builtin", []) if m.get("name")]


def _user():
    d = _read_json(USER_REG) or {}
    return [m for m in d.get("models", []) if m.get("name")]


def _hidden():
    """被用户删掉的内置模型：清单在代码里改不了，记一份 hidden 名单不再列出（重新添加同一地址即恢复）。"""
    d = _read_json(USER_REG) or {}
    return set(safe_name(x) for x in d.get("hidden", []))


def _write_user(models, hidden=None):
    os.makedirs(ASSETS, exist_ok=True)
    if hidden is None:
        hidden = _hidden()
    tmp = USER_REG + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump({"models": models, "hidden": sorted(set(hidden))}, f, ensure_ascii=False, indent=1)
    os.replace(tmp, USER_REG)


def _file_ok(path, size_mb=None):
    """模型文件是否完整。🔴 只看"大于 10MB"不够：中转 scp 被打断会留下一个 841MB 的半截文件（2026-09-07 实遇），
    清单里记了 size_mb 的按体积核对（容差 2%），没记的退回 10MB 下限。"""
    if not os.path.isfile(path):
        return False
    size = os.path.getsize(path)
    if size_mb:
        return abs(size / 1e6 - float(size_mb)) <= max(5.0, 0.02 * float(size_mb))
    return size > 10 * 1024 * 1024


def list_models():
    """内置 + 用户添加，带就绪状态。字段：name/label/author/res/size_mb/url/builtin/ready/preview。"""
    out, seen = [], set()
    hidden = _hidden()
    for src, builtin in ((_builtin(), True), (_user(), False)):
        for m in src:
            name = safe_name(m["name"])
            if name in seen or (builtin and name in hidden):
                continue
            seen.add(name)
            p = model_path(name)
            ready = _file_ok(p, m.get("size_mb"))
            pv = preview_file(name)
            out.append({"name": name, "label": m.get("label") or name, "author": m.get("author", ""),
                        "res": m.get("res"), "size_mb": m.get("size_mb") or (round(os.path.getsize(p) / 1e6) if ready else None),
                        "url": m.get("url"), "builtin": builtin, "ready": ready,
                        "preview": pv if os.path.isfile(os.path.join(ASSETS, pv)) else None,
                        "note": m.get("note", "")})
    return out


def find_model(name):
    name = safe_name(name)
    for m in list_models():
        if m["name"] == name:
            return m
    return None


def add_user_model(name, label="", url=None, size_mb=None):
    name = safe_name(name)
    with _LOCK:
        hidden = _hidden()
        if name in hidden and any(safe_name(m["name"]) == name for m in _builtin()):
            hidden.discard(name)                      # 重新添加一个被删过的内置模型 = 恢复它（内置条目自带地址）
            _write_user(_user(), hidden)
            return name
        models = [m for m in _user() if safe_name(m["name"]) != name]
        models.append({"name": name, "label": (label or name)[:40], "url": url, "size_mb": size_mb,
                       "added": time.strftime("%Y-%m-%d %H:%M:%S")})
        _write_user(models, hidden)
    return name


def remove_model(name):
    """删除模型：用户添加的从清单去掉；内置的进 hidden 名单不再列出（user 2026-09-07 要求内置也能删）。两者都删文件与预览图。"""
    name = safe_name(name)
    with _LOCK:
        hidden = _hidden()
        if any(safe_name(m["name"]) == name for m in _builtin()):
            hidden.add(name)
        models = [m for m in _user() if safe_name(m["name"]) != name]
        _write_user(models, hidden)
    for p in (model_path(name), model_path(name) + ".part", os.path.join(ASSETS, preview_file(name))):
        try:
            os.remove(p)
        except FileNotFoundError:
            pass


remove_user_model = remove_model                  # 旧名兼容


# ---------------------------------------------------------------- 下载
def download(url, path, progress=None, timeout=60):
    """流式下载到 path（.part 临时文件，成功后原子改名）。progress(done_bytes, total_bytes)。
    🔴 不做静默兜底：失败直接抛，让上层把原因显示给用户。"""
    os.makedirs(os.path.dirname(path), exist_ok=True)
    tmp = path + ".part"
    req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
    with urllib.request.urlopen(req, timeout=timeout) as r, open(tmp, "wb") as f:
        total = int(r.headers.get("Content-Length") or 0)
        done = 0
        last = 0.0
        while True:
            b = r.read(1 << 20)
            if not b:
                break
            f.write(b)
            done += len(b)
            if progress and time.time() - last > 0.5:
                progress(done, total)
                last = time.time()
    if os.path.getsize(tmp) < 10 * 1024 * 1024:
        os.remove(tmp)
        raise RuntimeError("下载得到的文件太小（%d 字节），不是有效的 .dfm" % os.path.getsize(tmp) if os.path.exists(tmp) else "下载得到的文件太小")
    os.replace(tmp, path)
    if progress:
        progress(done, total or done)
    return path


def ensure_model(name, progress=None):
    """确保 assets/<name>.dfm 在本地；没有就按清单里的 url 下载。返回路径。"""
    m = find_model(name)
    if m is None:
        raise FileNotFoundError("清单里没有这个模型: %s" % name)
    p = model_path(name)
    if m["ready"]:
        return p
    if not m.get("url"):
        raise FileNotFoundError("模型文件不在本机或不完整，且没有下载地址: %s" % name)
    if os.path.isfile(p):                          # 半截/损坏文件：删掉重下，不静默沿用
        _log("模型文件体积不对（%.0fMB，清单 %sMB），删掉重下: %s" % (os.path.getsize(p) / 1e6, m.get("size_mb"), name))
        os.remove(p)
    _log("下载模型", name, "←", m["url"])
    t = time.time()
    download(m["url"], p, progress)
    _log("下载完成 %s %.0fMB %.0fs" % (name, os.path.getsize(p) / 1e6, time.time() - t))
    return p


# ---------------------------------------------------------------- 换脸器
def _box_mask(size, blur=0.3):
    """FaceFusion face_masker.create_box_mask(crop, 0.3, (0,0,0,0))"""
    b = int(size * 0.5 * blur)
    a = max(b // 2, 1)
    m = np.ones((size, size), np.float32)
    m[:a, :] = 0; m[-a:, :] = 0; m[:, :a] = 0; m[:, -a:] = 0
    return cv2.GaussianBlur(m, (0, 0), b * 0.25) if b > 0 else m


def _compose(A, B):
    """两个 2×3 仿射的复合 A∘B（先 B 后 A），返回 2×3。"""
    A3 = np.vstack([A, [0, 0, 1]]).astype(np.float64)
    B3 = np.vstack([B, [0, 0, 1]]).astype(np.float64)
    return (A3 @ B3)[:2].astype(np.float32)


class DfmSwapper:
    """与 validated_impl.RoiSwapper 同契约的 .dfm 换脸器。

    · occluder / occ_every / _occ_cache：遮挡感知与 RoiSwapper 同源——xseg 在 arcface-128 对齐空间跑，
      缓存供修复器复用；再经复合仿射搬到 DFM 裁剪空间参与融合（两次对齐是不同模板，不能直接乘）。
    · color_match：复用 validated_impl 的 LAB gain/bias（口径=最终融合权重），隔 cm_every 帧重算 + EMA，
      与 RoiSwapper 一致（§10.3：统计口径必须等于合成口径）。
    · morph：AMP 模型的 morph_value（SAEHD 模型没有此输入，忽略）。
    · edge_converge：为兼容 pipeline 每帧挂载而保留，DFM 输出自带脸区 mask，边界由模型决定，不做亮度收敛。
    """

    def __init__(self, path, providers=None, occluder=None, color_match=True, morph=1.0):
        import onnxruntime as ort
        from insightface.utils import face_align
        from validated_impl import _apply_lab, _lab_gain_bias
        self._fa, self._gb, self._apply = face_align, _lab_gain_bias, _apply_lab
        providers = providers or ["CUDAExecutionProvider", "CPUExecutionProvider"]
        so = ort.SessionOptions()
        self.sess = ort.InferenceSession(path, so, providers=providers)
        if "CUDAExecutionProvider" not in self.sess.get_providers():
            _log("⚠️ DFM 未跑在 CUDA 上:", self.sess.get_providers())
        ins = {i.name: i.shape for i in self.sess.get_inputs()}
        if "in_face:0" not in ins:
            raise ValueError("不是 DeepFaceLive 格式的模型（缺 in_face:0 输入）: %s" % list(ins))
        self.R = int(ins["in_face:0"][1])
        self.has_morph = "morph_value:0" in ins
        self.out_names = [o.name for o in self.sess.get_outputs()]
        for need in ("out_face_mask:0", "out_celeb_face:0", "out_celeb_face_mask:0"):
            if need not in self.out_names:
                raise ValueError("模型缺输出 %s: %s" % (need, self.out_names))
        self.path, self.name = path, os.path.splitext(os.path.basename(path))[0]
        self.box = _box_mask(self.R)
        self.occluder, self.occ_every, self._occ_cache, self._occ_n, self._Ma = occluder, 1, None, 0, None
        self.color_match, self.morph = color_match, float(morph)
        self.local_color, self.local_color_strength = True, 0.65    # 局部光照迁移（colorfit），pipeline 每帧同步
        self.cm_every, self.cm_ema, self._cm, self._cm_n = 4, 0.75, None, 0
        self.edge_converge = 0.0
        self.pad = 12
        self.last_ms = 0.0
        self.input_size = (self.R, self.R)
        _log("DFM 就绪 %s res=%d morph=%s providers=%s" % (self.name, self.R, self.has_morph, self.sess.get_providers()[0]))

    def infer_crop(self, crop):
        """吃 R×R 的 BGR 裁剪，返回 (换好的脸 uint8 BGR, 融合 mask float[0,1])——都在 DFM 裁剪空间。"""
        x = cv2.addWeighted(crop, 1.75, cv2.GaussianBlur(crop, (0, 0), 2), -0.75, 0).astype(np.float32) / 255.0
        feed = {"in_face:0": x[None]}
        if self.has_morph:
            feed["morph_value:0"] = np.array([self.morph], np.float32)
        t = time.perf_counter()
        o = dict(zip(self.out_names, self.sess.run(None, feed)))
        self.last_ms = (time.perf_counter() - t) * 1000
        face = np.clip(o["out_celeb_face:0"][0] * 255.0, 0, 255).astype(np.uint8)
        m = np.minimum(np.squeeze(o["out_face_mask:0"]), np.squeeze(o["out_celeb_face_mask:0"])).astype(np.float32)
        m = cv2.GaussianBlur(cv2.erode(m, np.ones((3, 3), np.uint8)), (0, 0), 6.25).clip(0, 1)
        if face.shape[:2] != (self.R, self.R):
            face = cv2.resize(face, (self.R, self.R))
        if m.shape[:2] != (self.R, self.R):
            m = cv2.resize(m, (self.R, self.R))
        return face, np.minimum(self.box, m)

    def __call__(self, img, tgt_face):
        R = self.R
        kps = np.asarray(tgt_face.kps, np.float32)
        M, _ = cv2.estimateAffinePartial2D(kps, TEMPLATE * R, method=cv2.RANSAC, ransacReprojThreshold=100)
        if M is None:
            return img
        crop = cv2.warpAffine(img, M, (R, R), borderMode=cv2.BORDER_REPLICATE)
        face, mask = self.infer_crop(crop)

        # 遮挡感知：xseg 在 arcface-128 空间（与 RoiSwapper/修复器同一口径），复合仿射搬进 DFM 空间
        if self.occluder is not None:
            if self._occ_cache is None or self._occ_n % self.occ_every == 0:
                aimg, Ma = self._fa.norm_crop2(img, kps, 128)
                self._occ_cache = self.occluder.face_mask(aimg, out_size=(128, 128))
                self._Ma = Ma
            self._occ_n += 1
            if self._Ma is not None:
                T = _compose(M, cv2.invertAffineTransform(self._Ma))          # 128 空间 → 图像 → DFM 空间
                occ = cv2.warpAffine(self._occ_cache.astype(np.float32), T, (R, R), borderValue=1.0)
                mask = mask * occ.clip(0, 1)
        else:
            self._occ_cache = None

        # 肤色/亮度匹配：模型人物的肤色统计拉到目标当前光照下（口径=融合权重 mask）
        if self.color_match and float(mask.sum()) > 500:
            if self._cm is None or self._cm_n % self.cm_every == 0:
                new_cm = self._gb(face, crop, mask)
                if self._cm is None:
                    self._cm = new_cm
                else:
                    a = self.cm_ema
                    self._cm = (a * self._cm[0] + (1 - a) * new_cm[0], a * self._cm[1] + (1 - a) * new_cm[1])
            self._cm_n += 1
            face = self._apply(face, self._cm)
        # 局部光照迁移（§11.30）：把目标脸的明暗分布搬到模型脸上——全局匹配治不了"整张脸一块平光"的面具感
        if self.local_color and float(mask.sum()) > 500:
            face = local_color_transfer(crop, face, mask, float(self.local_color_strength))

        # ROI 回贴（§3：全帧回贴是 65ms 大坑）
        h, w = img.shape[:2]
        IM = cv2.invertAffineTransform(M)
        corners = np.array([[[0, 0], [R, 0], [R, R], [0, R]]], np.float32)
        d = cv2.transform(corners, IM).reshape(-1, 2)
        x0 = max(int(d[:, 0].min()) - self.pad, 0); x1 = min(int(d[:, 0].max()) + self.pad, w)
        y0 = max(int(d[:, 1].min()) - self.pad, 0); y1 = min(int(d[:, 1].max()) + self.pad, h)
        if x1 - x0 <= 2 or y1 - y0 <= 2:
            return img
        IMr = IM.copy(); IMr[0, 2] -= x0; IMr[1, 2] -= y0
        fr = cv2.warpAffine(face, IMr, (x1 - x0, y1 - y0), borderMode=cv2.BORDER_REPLICATE)
        mk = cv2.warpAffine(mask, IMr, (x1 - x0, y1 - y0), borderValue=0.0).clip(0, 1)[..., None]
        roi = img[y0:y1, x0:x1].astype(np.float32)
        img[y0:y1, x0:x1] = (mk * fr + (1 - mk) * roi).astype(np.uint8)
        return img


def make_preview(sw, app, sample_img, name):
    """用固定样本脸跑一次模型，存 256×256 对齐裁剪作缩略图（用户添加的模型没有现成预览）。"""
    from insightface.utils import face_align
    faces = app.get(sample_img)
    if not faces:
        raise RuntimeError("样本图里没检出人脸")
    f = max(faces, key=lambda q: (q.bbox[2] - q.bbox[0]) * (q.bbox[3] - q.bbox[1]))
    out = sw(sample_img.copy(), f)
    crop = face_align.norm_crop(out, f.kps, 256)
    ok, buf = cv2.imencode(".jpg", crop, [cv2.IMWRITE_JPEG_QUALITY, 90])
    if not ok:
        raise RuntimeError("预览图编码失败")
    with open(os.path.join(ASSETS, preview_file(name)), "wb") as fh:     # 字节写：路径含中文时 imwrite 静默失败（§5）
        fh.write(buf.tobytes())
    return preview_file(name)


# ---------------------------------------------------------------- CLI：预取内置模型（部署用）
if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser(description="预取内置 DFM 模型到 assets/")
    ap.add_argument("--fetch-all", action="store_true")
    ap.add_argument("--list", action="store_true")
    a = ap.parse_args()
    sys.stdout.reconfigure(encoding="utf-8")
    for m in list_models():
        print("%-28s %-22s res=%-4s %5s MB  ready=%s  builtin=%s" % (m["name"], m["label"], m["res"], m["size_mb"], m["ready"], m["builtin"]))
        if a.fetch_all and m["builtin"] and not m["ready"]:
            ensure_model(m["name"], progress=lambda d, t: print("\r  %5.1f%%" % (100.0 * d / max(t, 1)), end=""))
            print()
