"""接口4：用户面部特征分析（调用火山方舟 豆包视觉模型 doubao-seed-1-6-vision）。

算法来源：/home/xsl/fuyan（FaceArk.py）。worker 把图片以 base64 data URI 传给方舟
多模态模型，模型返回一大堆人脸特征 JSON；本模块解析后映射出接口4 的英文优先字段
（face_shape 等），并保留 doubao 返回的全部中文字段。

⚠️ 这是**唯一调外网云模型**的接口（其余接口全本地）。API Key 走配置/环境变量，不入 git。
"""
from __future__ import annotations

import base64
import json
import logging
import os

logger = logging.getLogger("hair.worker")

ARK_BASE_URL = os.getenv("ARK_BASE_URL", "https://ark.cn-beijing.volces.com/api/v3")
ARK_MODEL = os.getenv("ARK_MODEL", "doubao-seed-1-6-vision-250815")

# doubao 中文键 → 接口4 英文优先字段（仅保留这 6 项）
_KEY_MAP = {
    "脸型": "face_shape",
    "眉形": "eyebrow_shape",
    "面部年龄": "facial_age",
    "动静类型": "dynamic_static_type",
    "性别": "gender",
    "基因风格": "gene_style",
}

# 仅请求接口4 需要的 6 个字段（+「图片是否有人脸」用于 1001 判定，不进最终输出）
_PROMPT = (
    "分析一下图片告诉我以下特征，只要答案，格式为json字符串，"
    "图片是否有人脸（有人/没人） "
    "脸型（圆形脸/心形脸/菱形脸/鹅蛋脸/方形脸/长形脸/瓜子脸） 眉形 "
    "面部年龄（给出区间年龄） 动静类型（静态型/动态型） 性别（男/女） "
    "基因风格（戏剧型/睿智型/自然型/古典型/优雅型/浪漫型/前卫型/少女型/少年型）"
)

_client = None


def _load_api_key() -> str | None:
    """ARK_API_KEY 环境变量优先，否则读 worker_config.json / gateway/config.json 的 ark_api_key。"""
    key = os.getenv("ARK_API_KEY")
    if key:
        return key
    base = os.path.dirname(__file__)
    for cfg_name in ("worker_config.json", "gateway/config.json"):
        cfg = os.path.join(base, cfg_name)
        if os.path.isfile(cfg):
            try:
                with open(cfg, encoding="utf-8") as f:
                    v = json.load(f).get("ark_api_key")
                    if v:
                        return v
            except Exception as e:  # noqa: BLE001
                logger.warning("读取 %s ark_api_key 失败：%s", cfg_name, e)
    return None


def get_client():
    global _client
    if _client is None:
        from volcenginesdkarkruntime import Ark
        key = _load_api_key()
        if not key:
            raise RuntimeError("缺少火山方舟 API Key（设 ARK_API_KEY 或 worker_config.json.ark_api_key）")
        _client = Ark(base_url=ARK_BASE_URL, api_key=key)
    return _client


def _parse_json(text: str) -> dict:
    """去掉 ```json 包裹后解析。"""
    s = text.strip()
    if s.startswith("```"):
        s = s.strip("`")
        if s[:4].lower() == "json":
            s = s[4:]
    return json.loads(s.strip())


def _image_to_url(image_bytes: bytes = None, image_url: str = None) -> str:
    """优先用现成 URL；否则把字节转 base64 data URI（doubao 兼容）。"""
    if image_url:
        return image_url
    fmt = "png" if image_bytes[:8] == b"\x89PNG\r\n\x1a\n" else "jpeg"
    return f"data:image/{fmt};base64," + base64.b64encode(image_bytes).decode()


def analyze_features(image_bytes: bytes = None, image_url: str = None):
    """调 doubao 视觉模型分析人脸特征。

    Returns: dict —— 仅含接口4 的 6 个英文字段(face_shape/eyebrow_shape/facial_age/
      dynamic_static_type/gender/gene_style)；**无人脸返回 None**（调用方据此判 1001）。
    """
    url = _image_to_url(image_bytes, image_url)
    resp = get_client().chat.completions.create(
        model=ARK_MODEL,
        messages=[{
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": url}},
                {"type": "text", "text": _PROMPT},
            ],
        }],
        max_tokens=1024,        # 限制输出长度，模型秒回
        temperature=0,          # 固定输出，无随机采样，提速+结果稳定
        stream=False,           # 关闭流式，单次返回结果更快
        extra_body={
            "thinking": {
                "type": "disabled",  # 彻底关闭深度思考模式，提速50%+
            },
        },
    )
    text = resp.choices[0].message.content
    logger.info("doubao raw response (first 500 chars): %s", text[:500])
    try:
        raw = _parse_json(text)         # doubao 原始中文字段
    except (json.JSONDecodeError, ValueError) as e:
        logger.error("doubao 返回非 JSON，原文: %s", text[:1000])
        raise RuntimeError(f"豆包模型返回格式异常，无法解析为 JSON：{text[:200]}") from e
    if not has_face(raw):
        return None
    # 只保留 6 个英文字段（doubao 缺某字段则跳过）
    return {en: raw[zh] for zh, en in _KEY_MAP.items() if zh in raw}


def has_face(features: dict) -> bool:
    """据 doubao 的「图片是否有人脸」判断。"""
    v = features.get("图片是否有人脸") or features.get("是否有人") or ""
    return "没人" not in str(v) and "没有" not in str(v)


if __name__ == "__main__":
    import sys
    path = sys.argv[1] if len(sys.argv) > 1 else "tests/fixtures/frontal.jpg"
    with open(path, "rb") as f:
        feats = analyze_features(image_bytes=f.read())
    if feats is None:
        print("无人脸（1001）")
    else:
        print(json.dumps(feats, ensure_ascii=False, indent=2))