"""
把各数据集的人脸特征抽取一次并缓存为 JSON，供调参脚本反复使用。

MediaPipe 关键点检测是调参循环里唯一的耗时环节，缓存后调参可以秒级迭代。
"""

from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import Dict, List

import cv2

sys.path.insert(0, str(Path(__file__).resolve().parent))

from face_shape_classifier import (  # noqa: E402
    _get_face_mesh,
    extract_face_features,
)

ROOT = Path(__file__).resolve().parent
IMG_EXT = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}


def iter_images(root: Path) -> List[Path]:
    return sorted(p for p in root.rglob("*") if p.suffix.lower() in IMG_EXT)


def features_for(path: Path) -> Dict[str, float] | None:
    bgr = cv2.imread(str(path))
    if bgr is None:
        return None
    h, w = bgr.shape[:2]
    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    res = _get_face_mesh().process(rgb)
    if not res.multi_face_landmarks:
        return None
    return extract_face_features(res.multi_face_landmarks[0].landmark, image_size=(w, h))


def main() -> None:
    out_path = Path(sys.argv[1]) if len(sys.argv) > 1 else ROOT / "cache" / "features.json"
    out_path.parent.mkdir(parents=True, exist_ok=True)

    sources = {
        # 6 张带标注的基准图，文件名即期望脸型
        "benchmark": [p for p in ROOT.joinpath("test_img").glob("*.png")],
        "girl": iter_images(ROOT / "test_img" / "girl"),
        "man": iter_images(ROOT / "test_img" / "man"),
        "dataset": iter_images(ROOT / "test_img" / "脸型测试集合"),
    }

    records = []
    failed = 0
    for source, paths in sources.items():
        for i, path in enumerate(paths, 1):
            feats = features_for(path)
            if feats is None:
                failed += 1
                continue
            rel = path.relative_to(ROOT)
            records.append(
                {
                    "source": source,
                    "path": rel.as_posix(),
                    "file": path.name,
                    # dataset 的上级目录名即原始分组（弱标签，非可信真值）
                    "group": path.parent.name if source == "dataset" else source,
                    "expected": path.stem if source == "benchmark" else None,
                    "features": feats,
                }
            )
            if i % 50 == 0 or i == len(paths):
                print(f"[{source}] {i}/{len(paths)}", flush=True)

    out_path.write_text(json.dumps(records, ensure_ascii=False), encoding="utf-8")
    print(f"\n写入 {out_path}：{len(records)} 条，检测失败 {failed} 张")


if __name__ == "__main__":
    main()