#!/usr/bin/env python3
"""百度智能云人脸 1：N 搜索（独立脚本，不接入本仓库 FastAPI）。

对应官方文档：人脸 1：N 搜索 — https://cloud.baidu.com/doc/FACE/s/Gk37c1uzc
接口：POST https://aip.baidubce.com/rest/2.0/face/v3/search

前置条件（本脚本不负责「建库 / 注册人脸」）：
  - 在控制台创建应用并开通「人脸识别」相关接口权限；
  - 已使用人脸库管理 API 或控制台建立用户组，并向库中注册用户与人脸照片；
  - 否则搜索会失败（例如未找到匹配用户、人脸库为空等）。人脸库管理说明见产品文档「人脸库管理」章节。

配置从环境变量读取；启动时会从**仓库根目录**下的 `.env` 与**当前工作目录**下的 `.env` 加载（需已安装 `python-dotenv`，随 pydantic-settings 提供）。

主要环境变量（详见仓库 `.env.example` 中 Baidu Face 节）：
  BAIDU_FACE_APP_ID、BAIDU_FACE_API_KEY、BAIDU_FACE_SECRET_KEY（必填）
  BAIDU_FACE_GROUP_ID_LIST（与命令行 --groups 二选一；格式以百度人脸库文档为准，非法值由接口返回错误码）

用法示例（输入为**文件夹**，遍历其下所有支持的图片并打印识别日志）：

  uv run python scripts/baidu_face_1n_search.py /path/to/photos

支持格式：PNG、JPG、JPEG、BMP（单张 base64 建议 <2M，分辨率 <1920x1080，以官方文档为准）。
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from aip import AipFace
from dotenv import load_dotenv

_PROJECT_ROOT = Path(__file__).resolve().parent.parent
_IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".bmp"}


def _validate_group_id_list(s: str) -> None:
    """仅校验列表非空与数量上限；group_id 字符集等由百度接口校验。"""
    parts = [p.strip() for p in s.split(",") if p.strip()]
    if not parts:
        print("错误：group_id_list 解析后为空。", file=sys.stderr)
        sys.exit(2)
    if len(parts) > 10:
        print("错误：group_id 最多 10 个（逗号分隔）。", file=sys.stderr)
        sys.exit(2)


def _load_dotenv_files() -> None:
    load_dotenv(_PROJECT_ROOT / ".env")
    load_dotenv()


def _env(name: str) -> str:
    return (os.environ.get(name) or "").strip()


def _env_int(name: str, default: int) -> int:
    v = _env(name)
    if v.isdigit() or (v.startswith("-") and v[1:].isdigit()):
        return int(v)
    return default


def _face_client() -> AipFace:
    app_id = _env("BAIDU_FACE_APP_ID")
    api_key = _env("BAIDU_FACE_API_KEY")
    secret = _env("BAIDU_FACE_SECRET_KEY")
    if not app_id or not api_key or not secret:
        print(
            "错误：未配置百度人脸凭据。\n"
            "请在 `.env` 或环境中设置：BAIDU_FACE_APP_ID、BAIDU_FACE_API_KEY、"
            "BAIDU_FACE_SECRET_KEY\n"
            "（参考仓库 `.env.example` 中 Baidu Face 节；与语音 BAIDU_SPEECH_* 可为不同应用。）",
            file=sys.stderr,
        )
        sys.exit(2)
    client = AipFace(app_id, api_key, secret)
    conn_ms = _env("BAIDU_FACE_CONNECTION_TIMEOUT_MS")
    sock_ms = _env("BAIDU_FACE_SOCKET_TIMEOUT_MS")
    if conn_ms.isdigit():
        client.setConnectionTimeoutInMillis(int(conn_ms))
    if sock_ms.isdigit():
        client.setSocketTimeoutInMillis(int(sock_ms))
    return client


def _read_image_base64(path: Path) -> str:
    if not path.is_file():
        raise FileNotFoundError(str(path))
    raw = path.read_bytes()
    if not raw:
        raise ValueError("empty file")
    return base64.b64encode(raw).decode("ascii")


def _ts() -> str:
    return datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds")


def _list_image_files(folder: Path, *, recursive: bool) -> list[Path]:
    if not folder.is_dir():
        print(f"错误：不是有效文件夹：{folder}", file=sys.stderr)
        sys.exit(2)
    if recursive:
        out: list[Path] = []
        for p in folder.rglob("*"):
            if p.is_file() and p.suffix.lower() in _IMAGE_SUFFIXES:
                out.append(p)
    else:
        out = [
            p
            for p in folder.iterdir()
            if p.is_file() and p.suffix.lower() in _IMAGE_SUFFIXES
        ]
    return sorted(out, key=lambda p: p.name.lower())


def _search_options_from_env_and_args(args: argparse.Namespace) -> dict[str, Any]:
    qc = _env("BAIDU_FACE_QUALITY_CONTROL") or "NONE"
    lc = _env("BAIDU_FACE_LIVENESS_CONTROL") or "NONE"
    if args.quality_control is not None:
        qc = args.quality_control
    if args.liveness_control is not None:
        lc = args.liveness_control
    max_n = _env_int("BAIDU_FACE_MAX_USER_NUM", 1) if args.max_user_num is None else args.max_user_num
    match_th = _env_int("BAIDU_FACE_MATCH_THRESHOLD", 80) if args.match_threshold is None else args.match_threshold
    return {
        "max_user_num": max(1, min(50, max_n)),
        "match_threshold": max(0, min(100, match_th)),
        "quality_control": qc,
        "liveness_control": lc,
    }


def _resolve_group_id_list(args: argparse.Namespace) -> str:
    g = (args.group_id_list or "").strip() or _env("BAIDU_FACE_GROUP_ID_LIST")
    if not g:
        print(
            "错误：未指定人脸组。\n"
            "请设置环境变量 BAIDU_FACE_GROUP_ID_LIST，或传入命令行：--groups a,b",
            file=sys.stderr,
        )
        sys.exit(2)
    return g


def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="百度人脸 1：N 搜索：在指定人脸库组中，对文件夹内每张照片做相似度检索并打印识别日志。"
    )
    p.add_argument(
        "folder",
        type=Path,
        help="含照片的文件夹路径（仅处理 PNG/JPG/JPEG/BMP）",
    )
    p.add_argument(
        "--groups",
        dest="group_id_list",
        default=None,
        help="人脸组 id，逗号分隔，最多 10 个；未传时使用环境变量 BAIDU_FACE_GROUP_ID_LIST",
    )
    p.add_argument(
        "--max-user-num",
        type=int,
        default=None,
        help="覆盖环境变量 BAIDU_FACE_MAX_USER_NUM；返回前 N 个最相似用户（1–50）",
    )
    p.add_argument(
        "--match-threshold",
        type=int,
        default=None,
        help="覆盖环境变量 BAIDU_FACE_MATCH_THRESHOLD；0–100，默认 80",
    )
    p.add_argument(
        "--quality-control",
        choices=("NONE", "LOW", "NORMAL", "HIGH"),
        default=None,
        help="覆盖环境变量 BAIDU_FACE_QUALITY_CONTROL；默认 NONE",
    )
    p.add_argument(
        "--liveness-control",
        choices=("NONE", "LOW", "NORMAL", "HIGH"),
        default=None,
        help="覆盖环境变量 BAIDU_FACE_LIVENESS_CONTROL；默认 NONE",
    )
    p.add_argument(
        "--recursive",
        action="store_true",
        help="递归包含子目录中的图片",
    )
    p.add_argument(
        "--json",
        action="store_true",
        help="每张照片输出一行 JSON（file + API 原样响应），便于脚本解析",
    )
    return p.parse_args()


def main() -> None:
    _load_dotenv_files()
    args = _parse_args()
    folder = args.folder.resolve()
    group_id_list = _resolve_group_id_list(args)
    _validate_group_id_list(group_id_list)
    options = _search_options_from_env_and_args(args)

    _group_log = f"[{_ts()}] 使用 group_id_list={group_id_list!r}"
    if args.json:
        print(_group_log, file=sys.stderr)
    else:
        print(_group_log)

    files = _list_image_files(folder, recursive=args.recursive)
    if not files:
        print(
            f"[{_ts()}] 文件夹内未找到支持格式的图片：{folder}（{', '.join(sorted(_IMAGE_SUFFIXES))}；可加 --recursive）",
            file=sys.stderr,
        )
        sys.exit(2)

    client = _face_client()
    n = len(files)
    any_error = False

    for i, path in enumerate(files, start=1):
        rel = path.name
        try:
            b64 = _read_image_base64(path)
        except (OSError, ValueError) as e:
            any_error = True
            print(
                f"[{_ts()}] [{i}/{n}] 文件 {rel!r} 读取失败：{e}",
                file=sys.stderr,
            )
            continue

        resp = client.search(b64, "BASE64", group_id_list, options)

        if args.json:
            line = {
                "file": str(path),
                "relpath": rel,
                "index": i,
                "total": n,
                "response": resp,
            }
            print(json.dumps(line, ensure_ascii=False))
            if resp.get("error_code", -1) != 0:
                any_error = True
            continue

        err = resp.get("error_code")
        if err != 0:
            any_error = True
            msg = resp.get("error_msg", "")
            print(
                f"[{_ts()}] [{i}/{n}] {rel!r} 识别失败  error_code={err}  error_msg={msg!r}"
            )
            continue

        result = resp.get("result") or {}
        users = result.get("user_list") or []
        if not users:
            print(
                f"[{_ts()}] [{i}/{n}] {rel!r} 无匹配用户  user_list 为空（可检查人脸库或调低匹配阈值）"
            )
            continue

        for r, u in enumerate(users, start=1):
            gid = u.get("group_id", "")
            uid = u.get("user_id", "")
            info = u.get("user_info", "")
            score = u.get("score", "")
            tag = f"  [{r}]" if len(users) > 1 else ""
            print(
                f"[{_ts()}] [{i}/{n}] {rel!r} 识别成功{tag}  group_id={gid!r}  "
                f"user_id={uid!r}  user_info={info!r}  score={score}"
            )

    if any_error:
        sys.exit(1)


if __name__ == "__main__":
    main()