#!/usr/bin/env python3
"""
单人肖像：按目标岁数做年龄相貌推演（图生图）。
默认先用人脸检测（detect_and_embed）按与证件照相同规则做头肩裁切；多人同框时取检测框面积最大者。
可加 --skip-face-crop / --strict-face。与 Sophnet easyllms 同步 image-edit 交互。
"""

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

import requests
import sophnet_tools

from age_face_crop import prepare_age_portrait_crop

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-2509"
VALID_MODELS = [MODEL_DEFAULT]

DEFAULT_SIZE = "800*1200"
VALID_SIZES = ("800*1200", "960*1280", "1024*1024", "864*1184")

AGE_MIN = 1
AGE_MAX = 99

DEFAULT_NEGATIVE_PROMPT = (
    "换脸成他人，不像同一人，毁容级变形，身体比例失调，多张脸合照，水印，低俗，过度磨皮网红脸，纯卡通插画"
)


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


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 compose_age_prompt(age_years: int, extra_user_prompt: str) -> str:
    parts = [
        "人物肖像编辑：与参考图为同一人，保持身份可辨认，五官血缘与骨相特征连贯一致。",
        f"在自然真实的前提下，推演人物约为「{age_years}周岁」时的外貌——包括面部软组织、肤质细纹或稚气等与年龄相称的变化，不改变人种与性别。",
        "构图与取景可参考原图头肩比例；姿态与表情保持自然得体；光线柔和写实。",
        "禁止换脸为他人、禁止过度美颜滤镜、禁止纯卡通贴纸风；画面中仅一人。",
    ]
    if extra_user_prompt.strip():
        parts.append(f"用户补充：{extra_user_prompt.strip()}")
    return "".join(parts)


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 collect_target_ages(age_args, ages_csv):
    out = []
    if age_args:
        out.extend(int(a) for a in age_args)
    if ages_csv:
        for chunk in ages_csv.replace("，", ",").split(","):
            chunk = chunk.strip()
            if chunk:
                out.append(int(chunk))
    seen = set()
    uniq = []
    for a in out:
        if a not in seen:
            seen.add(a)
            uniq.append(a)
    return uniq


def validate_ages(ages):
    invalid = [a for a in ages if a < AGE_MIN or a > AGE_MAX]
    if invalid:
        print(
            f"Error: 年龄须在 {AGE_MIN}-{AGE_MAX} 之间整数岁，不合规：{invalid}",
            file=sys.stderr,
        )
        sys.exit(1)


def call_image_edit(
    api_key,
    prompt,
    model,
    image_refs,
    size=DEFAULT_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="age_look_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 emit_one_age_image(api_key, args, ages_list, idx, edit_refs):
    age = ages_list[idx]
    prompt = compose_age_prompt(age, args.prompt)
    result = call_image_edit(
        api_key,
        prompt,
        args.model,
        edit_refs,
        size=args.size,
        watermark=args.watermark,
        negative_prompt=args.negative_prompt,
        prompt_extend=args.prompt_extend,
    )
    raw_urls = extract_urls(result)
    if not raw_urls:
        print("STATUS=failed", flush=True)
        print(f"ERROR_MESSAGE=TARGET_AGE_{age}: url not found in API response.", flush=True)
        print(json.dumps(result, ensure_ascii=False), file=sys.stderr)
        sys.exit(1)

    raw_url = raw_urls[0]
    print(f"TARGET_AGE={age}", flush=True)
    signed_url, local_path = reupload_for_signed_url(api_key, raw_url)
    print(f"IMAGE_URL={signed_url or raw_url}", flush=True)
    if local_path:
        if signed_url:
            try:
                os.unlink(local_path)
            except OSError:
                pass
        else:
            print(f"PREVIEW_PATH={local_path}", flush=True)
    if idx + 1 < len(ages_list):
        time.sleep(max(0, args.pause_between_calls))


def main():
    parser = argparse.ArgumentParser(
        description="Sophnet portrait: simulate appearance at given age(s)",
    )
    parser.add_argument("--image", dest="images", action=AppendImageAction, required=True, help="Exactly one portrait")
    parser.add_argument(
        "--age",
        dest="ages",
        action="append",
        type=int,
        default=[],
        help="Target age in years (repeat for multiple)",
    )
    parser.add_argument(
        "--ages",
        dest="ages_csv",
        default=None,
        help='Comma-separated ages, e.g. "10,20,30"',
    )
    parser.add_argument("--prompt", default="", help="Optional user notes appended to edit instruction")
    parser.add_argument("--model", default=MODEL_DEFAULT, choices=VALID_MODELS)
    parser.add_argument(
        "--size",
        default=DEFAULT_SIZE,
        choices=VALID_SIZES,
        metavar="WxH",
        help=f"parameters.size (default {DEFAULT_SIZE})",
    )
    parser.add_argument("--watermark", type=parse_bool, default=False)
    parser.add_argument("--negative-prompt", default=DEFAULT_NEGATIVE_PROMPT)
    parser.add_argument("--prompt-extend", type=parse_bool, default=True)
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument(
        "--pause-between-calls",
        type=float,
        default=1.0,
        help="Seconds to sleep between sequential API calls (default 1.0)",
    )
    parser.add_argument(
        "--skip-face-crop",
        action="store_true",
        help="跳过人脸检测与头肩裁切，直接使用整张图送入年龄推演",
    )
    parser.add_argument(
        "--strict-face",
        action="store_true",
        help="必须通过人脸检测；未检出置信度达标人脸时不继续（退出码 2）",
    )
    args = parser.parse_args()

    ages_list = collect_target_ages(args.ages or [], args.ages_csv)

    if not ages_list:
        print("Error: specify at least one target age via --age or --ages.", file=sys.stderr)
        sys.exit(1)

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

    validate_ages(ages_list)

    ages_joined = ",".join(str(a) for a in ages_list)
    prompt_preview = compose_age_prompt(ages_list[0], args.prompt)

    if args.dry_run:
        print(f"MODEL={args.model}")
        print(f"SIZE={args.size}")
        print(f"INPUT_IMAGE_COUNT={len(args.images)}")
        print(f"TARGET_AGES={ages_joined}")
        fc = (
            "skipped_disabled"
            if args.skip_face_crop
            else ("strict_gate" if args.strict_face else "enabled_multi_largest_face")
        )
        print(f"FACE_CROP_PREP={fc}")
        print("STATUS=dry_run")
        brief = prompt_preview.replace("\n", " ")
        max_len = 400
        if len(brief) > max_len:
            brief = brief[:max_len] + "..."
        print(f"PROMPT_BRIEF={brief}")
        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_age_portrait_crop(
        args.images[0],
        api_key,
        args.size,
        skip_face_crop=args.skip_face_crop,
        strict_face_gate=args.strict_face,
    )
    if prep.get("error") == "no_face":
        print("STATUS=rejected", flush=True)
        print("REJECT_REASON=no_qualified_face", flush=True)
        print(
            "REJECT_HINT=未检出置信度达标的人脸，无法裁出单人头肩区。请换一张面部清晰的单人照或去掉 --strict-face。",
            flush=True,
        )
        sys.exit(2)

    edit_refs = [prep["edit_ref"]]
    cleanup_paths = prep.get("cleanup") or []

    print(f"MODEL={args.model}", flush=True)
    print(f"INPUT_IMAGE_COUNT={len(args.images)}", flush=True)
    print(f"TARGET_AGES={ages_joined}", flush=True)
    print(f"FACE_CROP={prep.get('face_crop', 'unknown')}", flush=True)
    n_face = prep.get("face_detected_count")
    if n_face is not None:
        print(f"FACE_DETECTED_COUNT={n_face}", flush=True)

    try:
        for i in range(len(ages_list)):
            emit_one_age_image(api_key, args, ages_list, i, edit_refs)
        print(f"OUTPUT_COUNT={len(ages_list)}", flush=True)
        print("STATUS=succeeded", flush=True)
    finally:
        for p in cleanup_paths:
            try:
                if p and os.path.isfile(p):
                    os.unlink(p)
            except OSError:
                pass


if __name__ == "__main__":
    main()
