yingqing/core/backend/apps/ai/mentions.py

"""全能创作 · @引用:实体检索 与 Ref 解析(契约 §1/§3)。

两件事:
1. `search_mentions()` —— 输入框打 @ 时的检索,返回 [Ref] 给前端渲染菜单。
2. `resolve_refs()` —— 把消息里的 [Ref] 变成模型真正吃得下的两样东西:
   **事实文本**(卖点/规格,进提示词)+ **参考图**(进 content_items,锁脸/锁商品/锁场景)。

铁律:消息里存的是结构化 Ref(type + id),**不是** "@净颜精华" 这串字。
后端必须拿 id 回表取事实与图,靠字符串匹配迟早对不上。
"""
from __future__ import annotations

import uuid
from dataclasses import dataclass, field

from django.db.models import Q

from apps.assets.models import Asset, Model
from apps.products.models import Product

from .services import _asset_preview_url, _product_cover_url

# Ref.type → 前端菜单里的分组名(设计稿 .omni-mention-group 的 small 文案)
TYPE_LABELS = {
    "product": "商品库",
    "model": "模特库",
    "character": "角色",
    "scene": "场景库",
    "asset": "资产库",
}
VALID_TYPES = tuple(TYPE_LABELS)
DEFAULT_TYPES = VALID_TYPES

# 参考图顺序固定:角色 → 场景 → 商品。这个顺序是出片模型 @图N 的语义依据,别改。
_REF_ORDER = {"model": 0, "character": 0, "scene": 1, "product": 2, "asset": 3}
# 火山单次出片最多 9 张图;留足余量,超出的靠优先级截断而不是报错
MAX_REFERENCE_IMAGES = 6

# 「资产库」是兜底分组,不重复列已经有专属分组的资产 ——
# 否则同一张定妆照会在「角色」和「资产库」各出现一次,菜单里看着像两个素材。
ASSET_EXCLUDED_CATEGORIES = (
    Asset.Category.PERSON,        # → character
    Asset.Category.SCENE,         # → scene
    Asset.Category.MODEL_PORTRAIT,  # → model
    Asset.Category.TRI_VIEW,        # → model
)


@dataclass
class ResolvedRefs:
    """resolve_refs 的产物。facts 进提示词,references 进 content_items。"""

    facts: list[str] = field(default_factory=list)
    references: list[dict] = field(default_factory=list)
    missing: list[dict] = field(default_factory=list)  # 删掉/不属于本团队的引用,要在对话里告诉用户

    @property
    def facts_text(self) -> str:
        return "\n\n".join(self.facts)


def _ref(type_: str, obj_id, name: str, cover: str = "") -> dict:
    return {"type": type_, "id": str(obj_id), "name": name, "cover": cover}


def _search_products(team, q: str, limit: int) -> list[dict]:
    queryset = Product.objects.filter(team=team, purged_at__isnull=True)
    if q:
        queryset = queryset.filter(title__icontains=q)
    out = []
    for product in queryset.order_by("-created_at", "-id")[:limit]:
        out.append(_ref("product", product.id, product.title, _product_cover_url(product)))
    return out


def _search_models(team, q: str, limit: int) -> list[dict]:
    queryset = Model.objects.filter(
        Q(team=team) | Q(is_official=True),
        is_deleted=False,
        purged_at__isnull=True,
    )
    if q:
        queryset = queryset.filter(name__icontains=q)
    out = []
    for model in queryset.select_related("portrait_asset").order_by("-is_official", "-created_at")[:limit]:
        out.append(_ref("model", model.id, model.name, _asset_preview_url(model.portrait_asset)))
    return out


def _search_assets(team, q: str, limit: int, categories: tuple[str, ...], type_: str) -> list[dict]:
    queryset = Asset.objects.filter(
        team=team,
        is_deleted=False,
        purged_at__isnull=True,
        asset_type=Asset.Type.IMAGE,
        category__in=categories,
    )
    if type_ == "asset":
        # 「资产库」只列用户真正加进库的图,不然工作台的每张试验图都会冒出来
        queryset = queryset.filter(in_library=True)
    if q:
        queryset = queryset.filter(name__icontains=q)
    out = []
    for asset in queryset.order_by("-created_at", "-id")[:limit]:
        out.append(_ref(type_, asset.id, asset.name, _asset_preview_url(asset)))
    return out


def search_mentions(team, q: str = "", types: list[str] | None = None, limit: int = 8) -> list[dict]:
    """@ 检索。types 不传则全类型各取 limit 条,按 商品 → 模特 → 角色 → 场景 → 资产 排。"""
    wanted = [t for t in (types or DEFAULT_TYPES) if t in TYPE_LABELS]
    q = (q or "").strip()
    results: list[dict] = []
    for type_ in wanted:
        if type_ == "product":
            results.extend(_search_products(team, q, limit))
        elif type_ == "model":
            results.extend(_search_models(team, q, limit))
        elif type_ == "character":
            results.extend(_search_assets(team, q, limit, (Asset.Category.PERSON,), "character"))
        elif type_ == "scene":
            results.extend(_search_assets(team, q, limit, (Asset.Category.SCENE,), "scene"))
        elif type_ == "asset":
            categories = tuple(
                c for c in Asset.Category.values if c not in ASSET_EXCLUDED_CATEGORIES
            )
            results.extend(_search_assets(team, q, limit, categories, "asset"))
    return results


_KEY_TO_TYPES = {
    "product": ["product"],
    "sku": ["product"],
    "goods": ["product"],
    "item": ["product"],
    "model": ["model"],
    "character": ["character"],
    "person": ["character"],
    "scene": ["scene"],
    "asset": ["asset"],
}


def infer_field_types(field: dict) -> list[str]:
    """追问卡字段 → 该去哪类库里解析用户的选择。"""
    typed = [t for t in (field.get("asset_types") or []) if t in TYPE_LABELS]
    if typed:
        return typed
    key = str(field.get("key") or "").strip().lower()
    if key in _KEY_TO_TYPES:
        return _KEY_TO_TYPES[key]
    label = str(field.get("label") or "")
    if "商品" in label:
        return ["product"]
    if "模特" in label:
        return ["model"]
    if "角色" in label or "人物" in label:
        return ["character"]
    if "场景" in label:
        return ["scene"]
    return list(DEFAULT_TYPES)


def lookup_mention(team, value: str, types: list[str] | None = None) -> dict | None:
    """把追问卡里的选项值(实体 id 或精确名字)还原成 Ref。对不上就返回 None,绝不瞎配。"""
    raw = str(value or "").strip()
    if not raw:
        return None
    wanted = [t for t in (types or DEFAULT_TYPES) if t in TYPE_LABELS] or list(DEFAULT_TYPES)
    uid = None
    try:
        uid = str(uuid.UUID(raw))
    except ValueError:
        uid = None
    if uid:
        if "product" in wanted:
            product = Product.objects.filter(team=team, id=uid, purged_at__isnull=True).first()
            if product:
                return _ref("product", product.id, product.title, _product_cover_url(product))
        if "model" in wanted:
            model = (
                Model.objects.filter(
                    Q(team=team) | Q(is_official=True),
                    id=uid,
                    is_deleted=False,
                    purged_at__isnull=True,
                )
                .select_related("portrait_asset")
                .first()
            )
            if model:
                return _ref("model", model.id, model.name, _asset_preview_url(model.portrait_asset))
        if any(t in wanted for t in ("character", "scene", "asset")):
            asset = Asset.objects.filter(
                team=team, id=uid, is_deleted=False, purged_at__isnull=True
            ).first()
            if asset is not None:
                if asset.category == Asset.Category.PERSON:
                    type_ = "character"
                elif asset.category == Asset.Category.SCENE:
                    type_ = "scene"
                else:
                    type_ = "asset"
                if type_ in wanted:
                    return _ref(type_, asset.id, asset.name, _asset_preview_url(asset))
        if types:
            return lookup_mention(team, raw, None)
        return None
    hits = search_mentions(team, q=raw, types=wanted, limit=8)
    for hit in hits:
        if (hit.get("name") or "") == raw:
            return hit
    return None


def refs_from_elicit_answers(team, fields, answers: dict) -> list[dict]:
    """用户在追问卡里点选的商品/角色等 → 可 pin 的 Ref。
    模型常用 type=single + 选项 value=商品id/名字,前端只回 answers 不回 refs,
    不在这里补上的话出片参考图里就没有这件商品。"""
    refs: list[dict] = []
    seen: set[tuple] = set()
    for field in fields or []:
        if not isinstance(field, dict):
            continue
        key = str(field.get("key") or "")
        if key in {"duration", "ratio", "resolution", "video_model", "count"}:
            continue
        raw = (answers or {}).get(field.get("key"))
        if raw is None:
            continue
        values = raw if isinstance(raw, list) else [raw]
        inferred = infer_field_types(field)
        for value in values:
            ref = lookup_mention(team, str(value or ""), inferred)
            if ref is None:
                continue
            mark = (ref.get("type"), str(ref.get("id")))
            if mark in seen:
                continue
            seen.add(mark)
            refs.append(ref)
    return refs


def _person_meta_facts(metadata) -> list[str]:
    """角色/模特库里填过的性别、年龄。聊天模型不一定看得到图,这些字必须进事实块。"""
    meta = metadata if isinstance(metadata, dict) else {}
    lines = []
    gender = str(meta.get("gender") or meta.get("sex") or "").strip()
    age = str(meta.get("age") or meta.get("age_range") or "").strip()
    if gender:
        lines.append(f"性别:{gender}")
    if age:
        lines.append(f"年龄:{age}")
    return lines


def product_facts_text(product) -> str:
    """商品事实块。全能创作没有 project,所以不能复用 script_agent._product_context()。
    这里只给**客观事实**(标题/品牌/品类/规格/卖点),不带人设和口吻 —— 那些由策略卡决定。"""
    lines = [f"商品:{product.title}"]
    if product.brand:
        lines.append(f"品牌:{product.brand}")
    if product.category:
        lines.append(f"品类:{product.category}")
    if product.target_audience:
        lines.append(f"目标人群:{product.target_audience}")
    description = (product.description or "").strip()
    if description:
        lines.append(f"商品描述:{description}")
    specs = product.specs if isinstance(product.specs, dict) else {}
    spec_text = "、".join(f"{k}:{v}" for k, v in specs.items() if v)
    if spec_text:
        lines.append(f"规格:{spec_text}")
    points = list(product.selling_points.order_by("sort_order", "created_at"))
    if points:
        joined = "\n".join(f"- {p.title}:{p.detail or p.title}" for p in points)
        lines.append(f"卖点:\n{joined}")
    return "\n".join(lines)


def _asset_reference(asset, type_: str, label: str) -> dict | None:
    """Asset → 参考图条目。带上审核态,视频路据此换成火山 asset:// 引用(否则真人图会被判「疑似真人」拒)。"""
    url = _asset_preview_url(asset)
    if not url:
        return None
    return {
        "url": url,
        "type": type_,
        "label": label,
        "source": "asset",
        "asset_id": str(asset.id),
        "review_status": asset.review_status,
        "review_remote_id": asset.review_remote_id,
    }



def _product_triview_asset(product):
    """商品三视图是独立 Asset(metadata.view=three_view),不在商品图列表里。没有就返回 None。"""
    return (
        Asset.objects.filter(
            team_id=product.team_id,
            is_deleted=False,
            purged_at__isnull=True,
            metadata__product_id=str(product.id),
            metadata__view="three_view",
        )
        .order_by("-created_at")
        .first()
    )


def _character_triview_asset(team, portrait_id):
    """角色立绘配对的三视图:metadata.triview_of = 立绘 id。没有就返回 None。"""
    if not portrait_id:
        return None
    return (
        Asset.objects.filter(
            team=team,
            is_deleted=False,
            purged_at__isnull=True,
            category=Asset.Category.TRI_VIEW,
            metadata__triview_of=str(portrait_id),
        )
        .order_by("-created_at")
        .first()
    ) or (
        Asset.objects.filter(
            team=team,
            is_deleted=False,
            purged_at__isnull=True,
            metadata__triview_of=str(portrait_id),
        )
        .order_by("-created_at")
        .first()
    )


def _newest_asset(*assets):
    """多张可用图里取最新一张。已删/空值直接跳过。"""
    living = [
        asset for asset in assets
        if asset is not None and not getattr(asset, "is_deleted", False)
    ]
    if not living:
        return None
    return max(living, key=lambda asset: (asset.created_at, str(asset.id)))


def _single_entity_asset(cover, *triviews):
    """一个角色/商品只喂一张参考图:有三视图用最新三视图,没有才用封面/立绘。

    封面和三视图衣服经常不一致;两张一起送,出片模型会当成多出来的角色或 SKU。
    """
    return _newest_asset(*triviews) or cover


def _append_ref(resolved: ResolvedRefs, asset, type_: str, label: str) -> bool:
    entry = _asset_reference(asset, type_, label)
    if not entry:
        return False
    resolved.references.append(entry)
    return True


def _product_reference(product) -> dict | None:
    """商品参考图:**真实上传图优先,排除 AI 生成图** —— 拿生成图当真相再喂回模型会误差累积。
    一张真实图都没有才回落封面(可能是 AI 图,但好过纯文生图)。

    这里不复用 services._product_reference_urls():那个只返回 url,而视频路还需要
    asset_id 和审核态才能把图换成火山 asset:// 引用(商品图也可能出现真人上身)。
    """
    rels = sorted(product.images.select_related("asset").all(), key=lambda im: (not im.is_primary, im.sort_order))
    for rel in rels:
        asset = rel.asset
        if asset is None or asset.source == Asset.Source.AI_GENERATED:
            continue
        entry = _asset_reference(asset, "product", product.title)
        if entry:
            return entry
    if product.cover_asset_id:
        entry = _asset_reference(product.cover_asset, "product", product.title)
        if entry:
            return entry
    cover = _product_cover_url(product)
    return (
        {"url": cover, "type": "product", "label": product.title,
         "source": "upload", "asset_id": "", "review_status": "", "review_remote_id": ""}
        if cover else None
    )


def resolve_refs(team, refs: list[dict]) -> ResolvedRefs:
    """[Ref] → 事实文本 + 参考图。查不到的进 missing,**不抛异常** ——
    素材被别人删掉不该让整条对话崩掉,该由 agent 在对话里说明。"""
    resolved = ResolvedRefs()
    for ref in refs or []:
        type_, ref_id = ref.get("type"), ref.get("id")
        if type_ not in TYPE_LABELS or not ref_id:
            continue
        if type_ == "product":
            product = Product.objects.filter(team=team, id=ref_id, purged_at__isnull=True).first()
            if product is None:
                resolved.missing.append(ref)
                continue
            resolved.facts.append(product_facts_text(product))
            chosen = _single_entity_asset(
                product.cover_asset,
                _product_triview_asset(product),
            )
            if chosen is not None and _append_ref(resolved, chosen, "product", product.title):
                continue
            entry = _product_reference(product)
            if entry:
                resolved.references.append(entry)
            elif chosen is None:
                resolved.missing.append(ref)
            continue
        if type_ == "model":
            model = Model.objects.filter(
                Q(team=team) | Q(is_official=True),
                id=ref_id,
                is_deleted=False,
                purged_at__isnull=True,
            ).select_related("triview_asset", "portrait_asset").first()
            if model is None:
                resolved.missing.append(ref)
                continue
            person_lines = []
            desc = (model.description or "").strip()
            meta_lines = _person_meta_facts(model.metadata)
            if desc:
                person_lines.append(f"模特「{model.name}」:{desc}")
            else:
                person_lines.append(f"模特「{model.name}」")
            person_lines.extend(meta_lines)
            if desc or meta_lines:
                resolved.facts.append("\n".join(person_lines))
            # 一人一张图:有三视图只用最新三视图。形象图和三视图衣服经常不一致,
            # 两张一起送会被当成多出来的角色。
            chosen = _single_entity_asset(
                model.portrait_asset,
                model.triview_asset,
                _character_triview_asset(team, getattr(model.portrait_asset, "id", None)),
            )
            if not _append_ref(resolved, chosen, "model", model.name):
                resolved.missing.append(ref)
            continue
        asset = Asset.objects.filter(
            team=team, id=ref_id, is_deleted=False, purged_at__isnull=True
        ).first()
        if asset is None:
            resolved.missing.append(ref)
            continue
        desc = (asset.description or "").strip()
        meta_lines = _person_meta_facts(asset.metadata)
        label = TYPE_LABELS[type_]
        person_lines = [f"{label}「{asset.name}」:{desc}" if desc else f"{label}「{asset.name}」"]
        person_lines.extend(meta_lines)
        if desc or meta_lines:
            resolved.facts.append("\n".join(person_lines))
        if type_ in {"character", "model", "asset"} and asset.category != Asset.Category.TRI_VIEW:
            chosen = _single_entity_asset(
                asset,
                _character_triview_asset(team, asset.id),
            )
        else:
            chosen = asset
        added = _append_ref(resolved, chosen, type_, asset.name)
        if not added:
            resolved.missing.append(ref)

    # 角色 → 场景 → 商品。同优先级内保持用户 @ 的先后。
    resolved.references.sort(key=lambda item: _REF_ORDER.get(item["type"], 9))
    resolved.references = _dedupe_references(resolved.references)[:MAX_REFERENCE_IMAGES]
    return resolved


def _dedupe_references(references: list[dict]) -> list[dict]:
    """同一张图被 @ 两次(比如商品图同时是资产库图)只留一条,否则 @图N 编号会错位。"""
    out: list[dict] = []
    seen: set[str] = set()
    for item in references:
        key = item.get("url") or ""
        if not key or key in seen:
            continue
        seen.add(key)
        out.append(item)
    return out