"""接口2 服务层：模型单例 + 性别贴图映射 + 「照片→N 张发际线预览图」管线。

把 head3d 的 extract_hairline 步骤包成单例复用（避免每请求重建模型），再按性别
对每张贴图调 render.render_hairline_overlay 生成预览图。
"""
from __future__ import annotations
import glob
import os

import cv2
import numpy as np

from . import constants as C
from .face_landmarks import FaceLandmarker
from .face_parsing import FaceParser
from .hairline_2d import sample_hairline, smooth_hairline
from .lift_3d import lift_hairline_to_3d, build_middle_row, assemble_full
from .render import load_ext_mesh, load_texture_rgba, render_hairline_overlay

_TEXTURE_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "hairline_texture")

# ⚠️ 本 worker 是 RTX 5090(sm_120)，torch 2.2.2(cu121) 只编到 sm_90，CUDA 跑算子会报
#    "no kernel image"。SegFormer 默认走 CPU（~2.5s/张）。换 torch cu128 后可设 SEG_DEVICE=cuda。
_SEG_DEVICE = os.getenv("SEG_DEVICE", "cpu")

_landmarker = None
_parser = None
_texture_map = None


def get_landmarker() -> FaceLandmarker:
    global _landmarker
    if _landmarker is None:
        _landmarker = FaceLandmarker(static_image_mode=True)
    return _landmarker


def get_parser() -> FaceParser:
    global _parser
    if _parser is None:
        _parser = FaceParser(device=_SEG_DEVICE)
    return _parser


def _gender_key(stem: str):
    """文件名 stem → (gender, key)；非 girl_/man_ 前缀返回 (None, None)。"""
    if stem.startswith("girl_"):
        return "female", stem[5:].replace(" ", "").strip()
    if stem.startswith("man_"):
        return "male", stem[4:].replace(" ", "").strip()
    return None, None


def get_texture_map() -> dict:
    """扫描 hairline_texture/ 建 {gender: [(key, path)]}，按 key 排序、缓存。

    文件名规范化去空格（如 `man_ inverse_arc.png` → key `inverse_arc`）。
    """
    global _texture_map
    if _texture_map is not None:
        return _texture_map
    mapping: dict[str, list] = {"female": [], "male": []}
    for path in sorted(glob.glob(os.path.join(_TEXTURE_DIR, "*.png"))):
        stem = os.path.splitext(os.path.basename(path))[0]
        gender, key = _gender_key(stem)
        if gender:
            mapping[gender].append((key, path))
    for g in mapping:
        mapping[g].sort(key=lambda kp: kp[0])
    _texture_map = mapping
    return _texture_map


def extract_502(image_bgr: np.ndarray):
    """照片(BGR) → (points502 MP序, valid17)。无人脸返回 (None, None)。"""
    rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
    landmarks = get_landmarker().detect(rgb)
    if landmarks is None:
        return None, None
    parse_map = get_parser().parse(rgb)
    hairline_2d, valid = sample_hairline(landmarks, parse_map)
    hairline_2d = smooth_hairline(hairline_2d, valid)
    hairline_3d = lift_hairline_to_3d(landmarks, hairline_2d)
    middle_3d = build_middle_row(landmarks, hairline_3d)
    points = assemble_full(landmarks, middle_3d, hairline_3d)
    return points, valid


def generate_previews(image_bgr: np.ndarray, gender: str):
    """生成该性别全部发际线预览图。

    Returns: list[dict]，每项 {"hairline_type": key, "image_bgr": ndarray, "order": 1..N}。
    无人脸返回 None。gender 必须是 male/female。
    """
    if gender not in ("male", "female"):
        raise ValueError(f"gender 必须是 male/female，收到 {gender!r}")
    points, _valid = extract_502(image_bgr)
    if points is None:
        return None

    uv, ext_faces = load_ext_mesh()
    results = []
    for order, (key, path) in enumerate(get_texture_map()[gender], start=1):
        tex = load_texture_rgba(path)
        preview = render_hairline_overlay(image_bgr, points, ext_faces, uv, tex)
        results.append({"hairline_type": key, "image_bgr": preview, "order": order})
    return results


if __name__ == "__main__":
    import sys
    g = sys.argv[2] if len(sys.argv) > 2 else "female"
    img = cv2.imread(sys.argv[1] if len(sys.argv) > 1 else "tests/fixtures/frontal.jpg")
    os.makedirs("tests/output", exist_ok=True)
    print("texture map:", {k: [kp[0] for kp in v] for k, v in get_texture_map().items()})
    res = generate_previews(img, g)
    if res is None:
        print("无人脸")
        sys.exit(1)
    for r in res:
        out = f"tests/output/preview_{g}_{r['hairline_type']}.png"
        cv2.imwrite(out, r["image_bgr"])
        print(f"  order={r['order']} type={r['hairline_type']} -> {out}")