cmshoppe/app/product_suite.py

"""Pure product-suite configuration and generation planning helpers."""

from __future__ import annotations

from collections import OrderedDict


FIXED_CATEGORIES = ("白底图", "场景图", "卖点图")
DEFAULT_CATEGORY_COUNTS = OrderedDict(
    (("白底图", 1), ("场景图", 2), ("卖点图", 2))
)
RATIOS = ("1:1", "3:4", "4:3", "16:9", "9:16")
MAX_CATEGORY_NAME_LENGTH = 10
MAX_GENERATION_COUNT_WITHOUT_CONFIRM = 16


def default_suite_settings():
    return {
        "platform": "Shopee",
        "country": "中国台湾",
        "language": "繁体中文",
        "ratio": "1:1",
        "per_image_primary": False,
        "categories": dict(DEFAULT_CATEGORY_COUNTS),
        "custom_category_order": [],
    }


def normalize_suite_settings(value=None):
    raw = dict(value or {}) if isinstance(value, dict) else {}
    normalized = default_suite_settings()
    normalized["platform"] = str(raw.get("platform") or "Shopee")
    normalized["country"] = str(raw.get("country") or "中国台湾")
    normalized["language"] = str(raw.get("language") or "繁体中文")
    ratio = str(raw.get("ratio") or "1:1")
    normalized["ratio"] = ratio if ratio in RATIOS else "1:1"
    normalized["per_image_primary"] = bool(raw.get("per_image_primary", False))

    raw_categories = raw.get("categories") if isinstance(raw.get("categories"), dict) else {}
    categories = OrderedDict()
    for name, default_count in DEFAULT_CATEGORY_COUNTS.items():
        categories[name] = _count(raw_categories.get(name, default_count))
    requested_order = raw.get("custom_category_order")
    requested_order = requested_order if isinstance(requested_order, list) else []
    seen = set(FIXED_CATEGORIES)
    custom_order = []
    for candidate in list(requested_order) + list(raw_categories):
        name = str(candidate or "")
        if name in seen or suite_name_error(name, categories):
            continue
        seen.add(name)
        custom_order.append(name)
        categories[name] = _count(raw_categories.get(name, 1))
    normalized["categories"] = dict(categories)
    normalized["custom_category_order"] = custom_order
    return normalized


def suite_name_error(name, existing=None, old_name=""):
    value = str(name or "")
    if not value.strip():
        return "分类名称不能为空"
    if value != value.strip() or any(character.isspace() for character in value):
        return "分类名称不能包含空格"
    if len(value) > MAX_CATEGORY_NAME_LENGTH:
        return "分类名称不能超过10个字"
    names = set(existing or ())
    if value in names and value != str(old_name or ""):
        return "分类名称已存在"
    return ""


def category_order(settings):
    normalized = normalize_suite_settings(settings)
    categories = normalized["categories"]
    custom = [
        name
        for name in normalized["custom_category_order"]
        if name in categories and name not in FIXED_CATEGORIES
    ]
    return list(FIXED_CATEGORIES) + custom


def suite_total_count(settings, image_count):
    normalized = normalize_suite_settings(settings)
    categories = normalized["categories"]
    base = sum(_count(categories.get(name, 0)) for name in category_order(normalized))
    if not normalized["per_image_primary"]:
        return base
    white_count = _count(categories.get("白底图", 0))
    other_count = max(0, base - white_count)
    return white_count + other_count * max(1, int(image_count or 0))


def build_suite_prompt(base_prompt, settings, category, item_id, source_index=1):
    normalized = normalize_suite_settings(settings)
    context = [
        "生成一张电商商品套图。",
        "平台：%s" % normalized["platform"],
        "国家地区：%s" % normalized["country"],
        "输出语言：%s" % normalized["language"],
        "图片比例：%s" % normalized["ratio"],
        "套图分类：%s" % str(category),
        "商品ID：%s" % str(item_id or ""),
        "当前主参考图序号：%d" % max(1, int(source_index or 1)),
        "商品卖点与要求：%s" % str(base_prompt or "").strip(),
        "保持商品主体、款式、颜色和关键细节准确，不添加无依据的功能或参数。",
    ]
    return "\n".join(context)


def build_job_specs(source_assets, base_prompt, settings, item_id):
    assets = list(source_assets or [])
    if not assets:
        return []
    normalized = normalize_suite_settings(settings)
    specs = []
    for category in category_order(normalized):
        count = _count(normalized["categories"].get(category, 0))
        if count <= 0:
            continue
        targets = assets if normalized["per_image_primary"] and category != "白底图" else assets[:1]
        for source_index, asset in enumerate(targets, 1):
            for category_index in range(1, count + 1):
                specs.append(
                    {
                        "source_asset_id": int(getattr(asset, "id", asset)),
                        "job_type": str(category),
                        "category": str(category),
                        "category_index": category_index,
                        "source_index": source_index,
                        "prompt": build_suite_prompt(
                            base_prompt,
                            normalized,
                            category,
                            item_id,
                            source_index=source_index,
                        ),
                    }
                )
    return specs


def _count(value):
    try:
        return max(0, int(value or 0))
    except (TypeError, ValueError):
        return 0