#!/usr/bin/env python3
"""
Sophnet 单图一键证件照：人脸检测（多人脸时取框面积最大者）→ 头肩裁切 → VL 预审（裁切图）→ 图改图（姿态校正、轻度遮挡整理与换底）。
`--skip-face-crop` 时跳过检测裁切，VL 仍按整张原图预审。输出含 FACE_CROP、COMPLIANCE_IMAGE_SCOPE、ID_PHOTO_BG、ALL_BG_TIP 等。
"""

import argparse
import base64
import json
import mimetypes
import os
import sys
import tempfile

import requests
import sophnet_tools

from id_photo_compliance import check_id_photo_compliance
from id_photo_face_crop import prepare_for_id_photo_edit

SOPH_BASE_URL = os.environ.get("SOPH_BASE_URL", "https://www.sophnet.com").rstrip("/")
API_URL = f"{SOPH_BASE_URL}/api/open-apis/projects/easyllms/image-edit"
MODEL_DEFAULT = "qwen-image-edit-plus"
VALID_MODELS = [MODEL_DEFAULT]

# 同步 image-edit 接口：parameters.size，竖版证件照常用
DEFAULT_ID_PHOTO_SIZE = "800*1200"
VALID_ID_PHOTO_SIZES = (DEFAULT_ID_PHOTO_SIZE, "960*1280")

ALL_BG_TIP = (
    "同一原图可再生成其它底色：加 --bg white（白底，默认）、--bg blue（蓝底）或 --bg red（红底）各执行一次即可。"
)

# 编辑提示：宜短；合规已拒口罩/墨镜等，此处仅指挥轻度遮挡与姿态。
ID_PHOTO_BASE = (
    "竖版证件照，同一人：换指定纯色底，轻调曝光与白平衡。"
    "姿态校正：身体直立、双肩略平，略正高低肩与头颈角度，面部朝正前；头肩取景、露双耳，头高约2/3、脸宽55%–60%、居中。"
    "可适当去除或拨开轻度遮挡：遮脸碎发、挡脸的手或误入镜小物，使五官与耳清晰；勿改脸型身份，光学眼镜保留。"
    "勿美颜、换脸、换装，勿去眼镜与水印。"
)

# bg id -> (中文名, 背景指令)
ID_BG = {
    "white": (
        "白底",
        "背景纯白平面，无渐变纹理。",
    ),
    "blue": (
        "蓝底",
        "背景标准证件蓝纯色平面，无渐变花纹。",
    ),
    "red": (
        "红底",
        "背景标准证件红纯色平面，无渐变花纹。",
    ),
}

MIME_MAP = {
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".png": "image/png",
    ".webp": "image/webp",
    ".gif": "image/gif",
    ".bmp": "image/bmp",
    ".avif": "image/avif",
}

# 与合规预审口径一致：关键词粗筛以合照、口罩、墨镜、人脸遮挡、耳朵为主（勿写场景特化词）。
VL_REJECT_RULES = (
    (
        "疑似合照，无法认定单一人物主体",
        (
            "多人",
            "多个人",
            "双人",
            "两人",
            "三人",
            "三名",
            "合照",
            "合影",
            "群照",
            "multiple people",
            "group photo",
            "two people",
            "three people",
        ),
    ),
    (
        "请勿佩戴口罩",
        (
            "戴口罩",
            "佩戴口罩",
            "戴着口罩",
            "wearing a mask",
        ),
    ),
    (
        "请勿佩戴墨镜",
        (
            "戴墨镜",
            "佩戴墨镜",
            "戴着墨镜",
            "wearing sunglasses",
        ),
    ),
    (
        "人脸被遮挡",
        (
            "遮挡面部",
            "遮住脸",
            "人脸被挡",
            "脸部被遮",
            "face is obscured",
        ),
    ),
    (
        "双耳未露出",
        (
            "双耳未露出",
            "耳朵未露出",
            "双耳不可见",
            "耳朵被头发遮挡",
            "遮住耳朵",
            "耳朵完全被挡",
        ),
    ),
)


def parse_bool(value):
    if value.lower() in ("true", "1", "yes", "y"):
        return True
    if value.lower() in ("false", "0", "no", "n"):
        return False
    raise argparse.ArgumentTypeError(f"invalid boolean value: {value!r}")


def is_remote_ref(ref):
    return ref.startswith(("http://", "https://", "data:image/"))


def path_to_data_uri(path):
    ext = os.path.splitext(path)[1].lower()
    mime = MIME_MAP.get(ext)
    if not mime:
        mime, _ = mimetypes.guess_type(path)
        if not mime:
            mime = "application/octet-stream"
    with open(path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode("ascii")
    return f"data:{mime};base64,{b64}"


def resolve_image(raw):
    ref = raw.lstrip("@")
    if is_remote_ref(ref):
        return ref
    if os.path.isfile(ref):
        return path_to_data_uri(ref)
    print(f"Error: --image input not found: {raw}", file=sys.stderr)
    print(
        "Hint: for uploaded files, use the resolved path from Media Understanding logs, "
        "usually under media/inbound/images/.",
        file=sys.stderr,
    )
    sys.exit(1)


def load_images_csv(csv_str):
    items = []
    for item in csv_str.split(","):
        item = item.strip()
        if item:
            items.append(item)
    return items


def load_images_file(path):
    if not os.path.isfile(path):
        print(f"Error: --images-file not found: {path}", file=sys.stderr)
        sys.exit(1)
    items = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip().rstrip("\r")
            if not line or line.startswith("#"):
                continue
            items.append(line)
    return items


def run_id_photo_compliance(image_ref: str, *, image_scope: str = "full"):
    """视觉模型判定输入图是否适合做证件照生成；image_scope cropped=裁切头肩区，full=整张原图。"""
    return check_id_photo_compliance(image_ref.lstrip("@"), image_scope=image_scope)


def load_text_file(path):
    if not os.path.isfile(path):
        print(f"Error: --vl-description-file (vision model caption file) not found: {path}", file=sys.stderr)
        sys.exit(1)
    with open(path, "r", encoding="utf-8") as f:
        return f.read().strip()


def keyword_is_negated(text, index):
    prefix = text[max(0, index - 12):index]
    return any(token in prefix for token in ("未", "没", "无", "不", "没有", "not ", "without "))


def has_unnegated_keyword(text, keyword):
    index = text.find(keyword)
    while index != -1:
        if not keyword_is_negated(text, index):
            return True
        index = text.find(keyword, index + len(keyword))
    return False


def reject_by_vl_description(description):
    text = " ".join(description.lower().split())
    if not text:
        return None
    for reason, keywords in VL_REJECT_RULES:
        if any(has_unnegated_keyword(text, keyword) for keyword in keywords):
            return reason
    return None


def compose_id_prompt(bg_id, extra_user_prompt):
    meta = ID_BG.get(bg_id)
    if not meta:
        valid = ", ".join(sorted(ID_BG))
        print(f"Error: unknown --bg {bg_id!r}. Valid: {valid}", file=sys.stderr)
        sys.exit(1)
    _label, bg_fragment = meta
    parts = [ID_PHOTO_BASE, bg_fragment]
    if extra_user_prompt:
        parts.append(f"用户补充：{extra_user_prompt.strip()}")
    return "".join(parts)


def call_image_edit(
    api_key,
    prompt,
    model,
    image_refs,
    size=DEFAULT_ID_PHOTO_SIZE,
    watermark=False,
    negative_prompt="",
    prompt_extend=True,
):
    resolved = [resolve_image(ref) for ref in image_refs]
    payload = {
        "model": model,
        "input": {
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"text": prompt},
                        *({"image": image} for image in resolved),
                    ],
                }
            ]
        },
        "parameters": {
            "negative_prompt": negative_prompt,
            "prompt_extend": prompt_extend,
            "watermark": watermark,
            "size": size,
        },
    }

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    resp = requests.post(API_URL, json=payload, headers=headers, timeout=120)
    resp.raise_for_status()
    return resp.json()


def extract_urls(data):
    urls = []

    def _walk(obj):
        if isinstance(obj, dict):
            for key in ("url", "image_url", "imageUrl"):
                value = obj.get(key)
                if isinstance(value, str) and value.startswith(("http://", "https://")):
                    urls.append(value)
            for v in obj.values():
                _walk(v)
        elif isinstance(obj, list):
            for item in obj:
                _walk(item)
        elif isinstance(obj, str) and obj.startswith(("http://", "https://")):
            urls.append(obj)

    _walk(data)
    return list(dict.fromkeys(urls))


def reupload_for_signed_url(api_key, raw_url):
    headers = {"Authorization": f"Bearer {api_key}"}
    try:
        resp = requests.get(raw_url, headers=headers, timeout=120, stream=True)
        resp.raise_for_status()
    except requests.RequestException as e:
        print(f"Warning: failed to download image: {e}", file=sys.stderr)
        return None, None

    ext = os.path.splitext(raw_url.split("?")[0])[-1] or ".png"
    fd, tmp_path = tempfile.mkstemp(suffix=ext, prefix="id_photo_reup_")
    try:
        with os.fdopen(fd, "wb") as f:
            for chunk in resp.iter_content(8192):
                f.write(chunk)
    except IOError as e:
        print(f"Warning: failed to write temp file: {e}", file=sys.stderr)
        os.unlink(tmp_path)
        return None, None

    signed_url = sophnet_tools.upload_oss(tmp_path)
    if not signed_url:
        print("Warning: upload_oss returned no signed URL, keeping local file", file=sys.stderr)
        return None, tmp_path

    return signed_url, tmp_path


def _unlink_temp_paths(paths):
    for tmp_path in paths or []:
        try:
            if tmp_path and os.path.isfile(tmp_path):
                os.unlink(tmp_path)
        except OSError:
            pass


class AppendImageAction(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        items = getattr(namespace, self.dest, None) or []
        items.append(values)
        setattr(namespace, self.dest, items)


def main():
    parser = argparse.ArgumentParser(
        description="Sophnet one-click ID photo from a single portrait",
    )
    parser.add_argument(
        "--bg",
        default="white",
        choices=sorted(ID_BG.keys()),
        help="Background: white (default), blue, red",
    )
    parser.add_argument("--prompt", default="", help="Extra notes only (e.g. crop); do not ask for beautification")
    parser.add_argument("--image", dest="images", action=AppendImageAction, default=[], help="Exactly one source image")
    parser.add_argument("--images", dest="images_csv", default=None, help="Exactly one image ref in this CSV")
    parser.add_argument("--images-file", default=None, help="Exactly one image path/URL per non-comment line")
    parser.add_argument(
        "--vl-description",
        default="",
        help="Optional caption from vision model (VL) for suitability filtering",
    )
    parser.add_argument(
        "--vl-description-file",
        default=None,
        help="UTF-8 text file: vision model (VL) caption for suitability filtering",
    )
    parser.add_argument("--model", default=MODEL_DEFAULT, choices=VALID_MODELS)
    parser.add_argument(
        "--size",
        default=DEFAULT_ID_PHOTO_SIZE,
        choices=VALID_ID_PHOTO_SIZES,
        metavar="WxH",
        help=f"parameters.size for image-edit (default {DEFAULT_ID_PHOTO_SIZE})",
    )
    parser.add_argument("--watermark", type=parse_bool, default=False)
    parser.add_argument("--negative-prompt", default="")
    parser.add_argument("--prompt-extend", type=parse_bool, default=True)
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument(
        "--skip-compliance-check",
        action="store_true",
        help="跳过视觉模型证件照输入合规预审（默认运行同目录 id_photo_compliance）",
    )
    parser.add_argument(
        "--skip-face-crop",
        action="store_true",
        help="跳过人脸检测与头肩裁切，直接将原图送图改图（慎用）",
    )
    parser.add_argument("--list-bg", action="store_true", help="Print --bg ids and labels, then exit")
    args = parser.parse_args()

    if args.list_bg:
        for bid in sorted(ID_BG.keys()):
            label, _ = ID_BG[bid]
            print(f"{bid}\t{label}")
        return

    all_images = list(args.images)
    if args.images_csv:
        all_images.extend(load_images_csv(args.images_csv))
    if args.images_file:
        all_images.extend(load_images_file(args.images_file))

    if not all_images:
        print("Error: at least one --image is required.", file=sys.stderr)
        sys.exit(1)

    if len(all_images) != 1:
        print(
            f"Error: 本工具仅支持 1 张输入图，当前为 {len(all_images)} 张。",
            file=sys.stderr,
        )
        sys.exit(1)

    vl_description = args.vl_description.strip()
    if args.vl_description_file:
        vl_description = load_text_file(args.vl_description_file)
    reject_reason = reject_by_vl_description(vl_description)
    if reject_reason:
        print("STATUS=rejected")
        print(f"REJECT_REASON={reject_reason}")
        print(
            "REJECT_HINT=请换一张单人照：勿合照；勿戴口罩或墨镜；人脸勿被遮挡；双耳需能露出。"
            f"（关键词命中：{reject_reason}）"
        )
        sys.exit(2)

    bg_label = ID_BG[args.bg][0]
    full_prompt = compose_id_prompt(args.bg, args.prompt)

    if args.dry_run:
        would_scope = "full" if args.skip_face_crop else "cropped"
        print(f"ID_PHOTO_BG={args.bg}")
        print(f"ID_PHOTO_BG_LABEL={bg_label}")
        print(f"ID_PHOTO_SIZE={args.size}")
        print(f"FACE_CROP={'skipped_disabled' if args.skip_face_crop else 'enabled'}")
        print(f"COMPLIANCE_IMAGE_SCOPE={would_scope}")
        print(f"ALL_BG_TIP={ALL_BG_TIP}")
        print(f"INPUT_IMAGE_COUNT={len(all_images)}")
        print("STATUS=dry_run")
        return

    api_key = sophnet_tools.get_api_key()
    if not api_key:
        print("Error: No API key found.", file=sys.stderr)
        sys.exit(1)

    prep = prepare_for_id_photo_edit(
        all_images[0],
        api_key,
        args.size,
        skip_face_crop=args.skip_face_crop,
        strict_face_gate=not args.skip_face_crop,
    )

    err = prep.get("error")
    if err == "no_face":
        print("STATUS=rejected", flush=True)
        print("REJECT_REASON=no_face", flush=True)
        print(
            "REJECT_HINT=未检测到足够置信度的人脸，请换一张面部清晰的近景照。",
            flush=True,
        )
        sys.exit(2)

    print(f"ID_PHOTO_BG={args.bg}", flush=True)
    print(f"ID_PHOTO_BG_LABEL={bg_label}", flush=True)
    print(f"ID_PHOTO_SIZE={args.size}", flush=True)
    print(f"FACE_CROP={prep.get('face_crop', 'unknown')}", flush=True)
    n_det = prep.get("face_detected_count")
    if n_det is not None and int(n_det) > 1:
        print(f"FACE_DETECTED_COUNT={int(n_det)}", flush=True)
        print("FACE_PICK=largest_area", flush=True)

    if not args.skip_compliance_check:
        scope = "cropped" if (not args.skip_face_crop and prep.get("face_crop") == "applied") else "full"
        try:
            report = run_id_photo_compliance(prep["edit_ref"], image_scope=scope)
        except Exception as e:
            print(f"Error: 证件照合规预审失败: {e}", file=sys.stderr)
            _unlink_temp_paths(prep.get("cleanup"))
            sys.exit(1)
        print(f"COMPLIANCE_IMAGE_SCOPE={scope}", flush=True)
        if not report.get("compliant"):
            _unlink_temp_paths(prep.get("cleanup"))
            print("STATUS=rejected")
            print("REJECT_REASON=compliance_vlm")
            print(f"COMPLIANCE_SUMMARY={report.get('summary', '')}")
            viols = report.get("violations") or []
            print(f"COMPLIANCE_VIOLATIONS={json.dumps(viols, ensure_ascii=False)}")
            print(f"COMPLIANCE_STANDARD_BRIEF={report.get('standard_brief', '')}")
            hint_parts = [str(report.get("summary", "")).strip()]
            if viols:
                hint_parts.append("不符合项：" + "；".join(str(v) for v in viols))
            hint_parts.append("完整规范说明见 COMPLIANCE_STANDARD_BRIEF。")
            print(f"REJECT_HINT={' '.join(p for p in hint_parts if p)}")
            sys.exit(2)

    try:
        result = call_image_edit(
            api_key,
            full_prompt,
            args.model,
            [prep["edit_ref"]],
            size=args.size,
            watermark=args.watermark,
            negative_prompt=args.negative_prompt,
            prompt_extend=args.prompt_extend,
        )
    finally:
        _unlink_temp_paths(prep.get("cleanup"))
    print(f"INPUT_IMAGE_COUNT={len(all_images)}")
    print("STATUS=succeeded")

    raw_urls = extract_urls(result)
    if not raw_urls:
        print("Error: url not found in response.", file=sys.stderr)
        print(json.dumps(result, ensure_ascii=False), file=sys.stderr)
        sys.exit(1)

    print(f"OUTPUT_COUNT={len(raw_urls)}")
    for i, raw_url in enumerate(raw_urls, 1):
        signed_url, local_path = reupload_for_signed_url(api_key, raw_url)
        print(f"IMAGE_URL={signed_url or raw_url}")
        if local_path:
            if signed_url:
                os.unlink(local_path)
            else:
                print(f"PREVIEW_PATH={local_path}")

    print(f"ALL_BG_TIP={ALL_BG_TIP}", flush=True)


if __name__ == "__main__":
    main()
