#!/usr/bin/env python3
"""
Sophnet 表情包图生图：固定模型 Qwen-Image-Edit-2509，仅支持恰好一张输入图；
预设风格随机/指名/完全自定义。输出关键字段与 sophnet-image-edit 对齐，另含 SELECTED_STYLE_*。
"""

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

import requests
import sophnet_tools

API_URL = "https://www.sophnet.com/api/open-apis/projects/easyllms/imagegenerator/task"
MODEL_DEFAULT = "Qwen-Image-Edit-2509"
VALID_MODELS = [MODEL_DEFAULT]

STICKER_STYLE_BASE = (
    "以输入图为准，保留原图主体、构图、表情和口型；"
    "只做轻量表情包贴纸化处理，画风不要覆盖原图识别度；"
    "背景简洁，无文字水印。"
)

# id 稳定，供 --style 与 --avoid-style 使用
STICKER_STYLES = [
    {"id": "kawaii_flat", "name": "日系Q版扁平贴纸", "prompt": "日系Q版扁平贴纸风，粗黑描边，高饱和可爱配色，哑光平涂"},
    {"id": "line_meme", "name": "极简线稿meme", "prompt": "极简黑白线稿meme风，白底，超粗轮廓线，夸张简笔五官"},
    {"id": "pixel_retro", "name": "复古8-bit像素", "prompt": "复古8-bit像素风表情包，有限色盘，马赛克边缘，怀旧游戏机感"},
    {"id": "watercolor_soft", "name": "柔边水彩", "prompt": "柔边水彩插画，纸纹轻微，治愈低对比配色，晕染边缘"},
    {"id": "clay_3d", "name": "3D黏土软陶", "prompt": "3D黏土/软陶质感，圆润体积，柔和明暗，盲盒手办感"},
    {"id": "american_toon", "name": "美式夸张卡通", "prompt": "美式卡通夸张动力学，粗轮廓，喜剧化五官，明亮色块"},
    {"id": "korean_sticker", "name": "韩系贴图闪粉", "prompt": "韩系通讯贴图画风，粉嫩配色，腮红高光，少量星星闪粉装饰不挡脸"},
    {"id": "vintage_sticker", "name": "轻复古贴纸", "prompt": "轻复古贴纸风，保留原图主色与轮廓，少量暖色滤镜和细小网点作为装饰"},
    {"id": "neon_outline", "name": "霓虹描边贴纸", "prompt": "保留原图主体与主色，仅在外轮廓添加柔和霓虹描边和少量发光点缀"},
    {"id": "ink_wash_light", "name": "淡彩水墨贴纸", "prompt": "淡彩水墨贴纸风，保留原图五官和服饰结构，仅加入轻微宣纸纹理与柔和墨边"},
    {"id": "retro_cartoon_soft", "name": "复古软卡通", "prompt": "复古软卡通贴纸风，圆润线条和轻微胶片色调，不改变角色比例与关键特征"},
    {"id": "pastel_blob", "name": "马卡龙软体", "prompt": "马卡龙pastel软体萌物，无锐利尖角，圆润渐变blob"},
    {"id": "sketch_pencil", "name": "铅笔速写", "prompt": "铅笔速写排线与纸张颗粒，未完成感俏皮速写"},
    {"id": "y2k_chrome", "name": "Y2K镭射金属", "prompt": "Y2K镭射金属贴纸，渐变幻彩，星星蝴蝶小装饰（不遮挡主体）"},
    {"id": "chalk_doodle_light", "name": "粉笔涂鸦边框", "prompt": "保留原图主体，仅添加轻微粉笔线条和手绘涂鸦边框，背景可简洁浅色"},
    {"id": "soft_felt_outline", "name": "柔软布贴边", "prompt": "柔软布贴贴纸风，只强化毛绒边缘和缝线点缀，保留原图色彩和五官细节"},
    {"id": "glass_flat", "name": "玻璃拟态扁平", "prompt": "玻璃拟态扁平插画，半透明磨砂块与柔和渐变，现代UI贴纸"},
    {"id": "red_stamp_frame", "name": "红章边框贴纸", "prompt": "保留原图完整彩色主体，仅添加红色印章感边框和少量中式装饰纹理"},
    {"id": "comic_action_light", "name": "轻漫画动感", "prompt": "轻漫画动感贴纸风，保留原图色彩与脸部细节，只添加少量速度线和表情强化"},
    {"id": "doodle_marker", "name": "粗头马克笔涂鸦", "prompt": "粗头马克笔儿童涂鸦，手绘颤抖线，高能量搞笑"},
    {"id": "clean_vector", "name": "干净矢量贴纸", "prompt": "干净矢量贴纸风，平滑色块和清晰描边，尽量沿用原图主色和外形比例"},
    {"id": "sparkle_outline", "name": "闪光描边贴纸", "prompt": "保留原图主体，仅添加透明贴纸边、细闪高光和少量星星装饰，不改变材质"},
    {"id": "soft_grain", "name": "柔颗粒插画", "prompt": "轻微柔颗粒插画质感，保留原图颜色关系与五官细节，只让边缘更贴纸化"},
    {"id": "plush_soft", "name": "轻毛绒贴纸", "prompt": "轻毛绒贴纸风，只增加柔软边缘和轻微绒感，不把主体改成玩偶或织物"},
    {"id": "emoji_soft3d", "name": "圆润软3D梗图", "prompt": "通讯软件常见圆润软体3D表情风，亚克力高光，软塑胶体积，简净场景"},
    {"id": "holo_edge", "name": "全息边缘贴纸", "prompt": "保留原图主体，仅在贴纸边缘加入轻微全息彩虹反光，内部颜色不大改"},
    {"id": "paper_texture_light", "name": "轻纸感贴纸", "prompt": "轻纸张纹理贴纸风，保留原图细节，只加入轻微纸纹、阴影和边缘厚度"},
    {"id": "crayon_soft", "name": "柔和蜡笔贴纸", "prompt": "柔和蜡笔贴纸风，保留原图轮廓与主色，轻微手绘颗粒，不做厚重覆盖"},
    {"id": "cel_anime_soft", "name": "赛璐璐柔光动画", "prompt": "赛璐璐动画赛璐片上色，柔和边光与空气透视，清爽线稿，日系剧场感"},
    {"id": "pastel_gradient", "name": "粉彩渐变贴纸", "prompt": "保留原图主体，加入柔和粉彩渐变背景和轻微发光边，主体颜色不大幅偏移"},
    {"id": "bold_outline", "name": "粗描边贴纸", "prompt": "只加强主体外轮廓和局部表情线条，保持原图内部色彩、服饰与五官识别度"},
    {"id": "japanese_pattern_light", "name": "和风轻纹样", "prompt": "保留原图主体，仅在背景或边框加入轻量和风纹样，不改变主体为版画"},
    {"id": "toy_photo_soft", "name": "软玩具摄影感", "prompt": "保留原图造型，轻微增加玩具般柔和光照和圆润质感，不改变为微缩场景"},
    {"id": "badge_outline", "name": "徽章描边贴纸", "prompt": "保留原图彩色主体，仅添加徽章式外框、白边和轻微阴影，不做金属化重绘"},
    {"id": "street_doodle_light", "name": "街头涂鸦点缀", "prompt": "保留原图主体，仅添加少量涂鸦线条、贴纸标签和喷点装饰，不覆盖面部"},
    {"id": "soft_print_grain", "name": "轻印刷颗粒", "prompt": "轻微印刷颗粒和柔和套色感，保留原图主色与细节，不使用强错位或有限色重绘"},
    {"id": "tiny_props", "name": "小道具贴纸", "prompt": "保留原图人物或主体，仅添加少量可爱小道具和表情符号，不改变主体材质"},
    {"id": "glossy_sticker", "name": "亮面贴纸", "prompt": "亮面贴纸质感，保留原图色彩与结构，仅增加透明高光、白边和柔和投影"},
    {"id": "rounded_app_icon", "name": "圆角图标贴纸", "prompt": "保留原图主体，整理成圆角图标式贴纸构图，背景简洁，不改变主体特征"},
    {"id": "cute_dark", "name": "暗色可爱贴纸", "prompt": "暗色系可爱贴纸风，保留原图主体颜色和五官，背景用深色柔和渐变与小装饰"},
]

_STYLE_BY_ID = {s["id"]: s for s in STICKER_STYLES}

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 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 compose_sticker_prompt(style_fragment, extra_user_prompt):
    parts = [STICKER_STYLE_BASE, f"画风与质感要求：{style_fragment}。"]
    if extra_user_prompt:
        parts.append(f"用户补充：{extra_user_prompt.strip()}")
    return "".join(parts)


def pick_style(random_mode, style_id, style_custom, avoid_ids):
    if sum(bool(x) for x in (random_mode, style_id, style_custom)) > 1:
        print("Error: use only one of --random-style, --style, --style-custom", file=sys.stderr)
        sys.exit(1)
    if not random_mode and not style_id and not style_custom:
        print("Error: specify one of --random-style, --style <id>, or --style-custom \"...\"", file=sys.stderr)
        sys.exit(1)

    if style_custom:
        return {
            "id": "custom",
            "name": "自定义",
            "prompt": style_custom.strip(),
        }

    if style_id:
        found = _STYLE_BY_ID.get(style_id)
        if not found:
            valid = ", ".join(sorted(_STYLE_BY_ID))
            print(f"Error: unknown --style {style_id!r}. Valid ids: {valid}", file=sys.stderr)
            sys.exit(1)
        return dict(found)

    avoid = set(avoid_ids or [])
    pool = [s for s in STICKER_STYLES if s["id"] not in avoid]
    if not pool:
        pool = list(STICKER_STYLES)
    return dict(random.choice(pool))


def create_task(api_key, prompt, model, image_refs,
                size="1024*1024", n=1, watermark=False):
    resolved = [resolve_image(ref) for ref in image_refs]
    payload = {
        "model": model,
        "input": {"prompt": prompt, "images": resolved},
        "parameters": {
            "size": size,
            "n": n,
            "watermark": watermark,
        },
    }

    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()
    data = resp.json()

    for key in ("task_id", "taskId", "taskID", "id"):
        if key in data:
            return data[key]
        if isinstance(data.get("output"), dict) and key in data["output"]:
            return data["output"][key]

    print("Error: task_id not found in response.", file=sys.stderr)
    print(json.dumps(data, ensure_ascii=False), file=sys.stderr)
    sys.exit(1)


def poll_task(api_key, task_id, poll_interval=2, max_wait=300):
    headers = {"Authorization": f"Bearer {api_key}"}
    start = time.time()

    while True:
        elapsed = time.time() - start
        if elapsed > max_wait:
            print(f"Error: timed out after {max_wait}s.", file=sys.stderr)
            sys.exit(1)

        resp = requests.get(f"{API_URL}/{task_id}", headers=headers, timeout=60)
        resp.raise_for_status()
        data = resp.json()

        status = None
        for key in ("status", "taskStatus", "task_status"):
            if key in data:
                status = data[key]
                break
            if isinstance(data.get("output"), dict) and key in data["output"]:
                status = data["output"][key]
                break

        if status and status.lower() in ("succeeded", "success"):
            return data
        if status and status.lower() in ("failed", "error"):
            print("STATUS=failed", file=sys.stderr)
            print(json.dumps(data, ensure_ascii=False), file=sys.stderr)
            sys.exit(1)

        time.sleep(poll_interval)


def extract_urls(data):
    urls = []

    def _walk(obj):
        if isinstance(obj, dict):
            if "url" in obj and isinstance(obj["url"], str):
                urls.append(obj["url"])
            for v in obj.values():
                _walk(v)
        elif isinstance(obj, list):
            for item in obj:
                _walk(item)

    _walk(data)
    return 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="sticker_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


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)


class AppendAvoidAction(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 sticker/emoji image edit with Qwen-Image-Edit-2509",
    )
    parser.add_argument("--prompt", default="", help="Optional extra user instructions appended to sticker prompt")
    parser.add_argument("--random-style", action="store_true", help="Randomly pick one preset sticker style")
    parser.add_argument("--style", default=None, metavar="ID", help="Preset style id (see --list-styles)")
    parser.add_argument("--style-custom", default=None, help="Fully custom style description (overrides presets)")
    parser.add_argument("--avoid-style", dest="avoid_styles", action=AppendAvoidAction, default=[],
                        help="When using --random-style, exclude this style id (repeatable)")
    parser.add_argument("--list-styles", action="store_true", help="Print preset ids and exit")
    parser.add_argument("--image", dest="images", action=AppendImageAction, default=[],
                        help="Exactly one source image (URL, data URI, or local path).")
    parser.add_argument("--images", dest="images_csv", default=None,
                        help="Exactly one image ref in this CSV (unsupported: multiple).")
    parser.add_argument("--images-file", default=None, help="Exactly one image path/URL in file")
    parser.add_argument("--model", default=MODEL_DEFAULT, choices=VALID_MODELS)
    parser.add_argument("--size", default="1024*1024")
    parser.add_argument("--n", type=int, default=1)
    parser.add_argument("--watermark", type=parse_bool, default=False)
    parser.add_argument("--poll-interval", type=int, default=2)
    parser.add_argument("--max-wait", type=int, default=300)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    if args.list_styles:
        for s in STICKER_STYLES:
            print(f"{s['id']}\t{s['name']}")
        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)} 张。多图合成请改用 sophnet-image-edit。",
            file=sys.stderr,
        )
        sys.exit(1)

    chosen = pick_style(
        args.random_style,
        args.style,
        args.style_custom,
        args.avoid_styles,
    )
    full_prompt = compose_sticker_prompt(chosen["prompt"], args.prompt)

    if args.dry_run:
        print(f"SELECTED_STYLE_ID={chosen['id']}")
        print(f"SELECTED_STYLE_NAME={chosen['name']}")
        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)

    print(f"SELECTED_STYLE_ID={chosen['id']}", flush=True)
    print(f"SELECTED_STYLE_NAME={chosen['name']}", flush=True)

    task_id = create_task(
        api_key, full_prompt, args.model, all_images,
        size=args.size, n=args.n, watermark=args.watermark,
    )
    print(f"INPUT_IMAGE_COUNT={len(all_images)}")
    print(f"TASK_ID={task_id}")

    result = poll_task(api_key, task_id, args.poll_interval, args.max_wait)

    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}")


if __name__ == "__main__":
    main()
