SupplyAgent/agents/generate_demand_agent/tools/dim_constants.py

"""generate_demand_agent 分类树热度维度共享常量与工具函数。"""
from __future__ import annotations

from collections.abc import Callable
from decimal import Decimal

from supply_infra.db.models.category_tree_weight import CategoryTreeWeight
from supply_infra.db.models.global_tree_category import GlobalTreeCategory

DIM_KEYS: tuple[str, ...] = (
    "ext_pop",
    "plat_sust_pop",
    "plat_ly_pop",
    "recent_pop",
)

DIM_LABEL: dict[str, str] = {
    "ext_pop": "外部热度",
    "plat_sust_pop": "平台持续热度",
    "plat_ly_pop": "平台去年同期热度",
    "recent_pop": "近期热度",
}

# 节点下存在有数据的叶子节点时的标记
HAS_DATA_LEAF_MARK = "+"


def normalize_biz_dt(biz_dt: str | None) -> tuple[str | None, str | None]:
    """校验并规范化 biz_dt（YYYYMMDD）；空值表示使用最新业务日。"""
    if biz_dt is None:
        return None, None
    text = str(biz_dt).strip()
    if not text:
        return None, None
    if len(text) != 8 or not text.isdigit():
        return None, f"biz_dt 格式无效，应为 YYYYMMDD: {biz_dt!r}"
    return text, None


def resolve_biz_dt(
    biz_dt: str | None,
    *,
    get_latest: Callable[[], str | None],
    has_data: Callable[[str], bool] | None = None,
    table_label: str = "热度数据",
) -> tuple[str | None, str | None]:
    """解析 biz_dt：未传则用最新业务日；传入则校验格式与数据是否存在。"""
    normalized, err = normalize_biz_dt(biz_dt)
    if err:
        return None, err

    resolved = normalized or get_latest()
    if not resolved:
        return None, f"暂无 {table_label}（请指定 biz_dt 或先写入数据）"

    if normalized and has_data is not None and not has_data(normalized):
        return None, f"biz_dt={normalized} 无 {table_label}"

    return resolved, None


def get_latest_common_biz_dt(session) -> str | None:
    """返回 category_tree_weight 与 demand_popularity_stats 均有数据的最新 biz_dt。"""
    from sqlalchemy import func, select

    from supply_infra.db.models.demand_popularity_stats import DemandPopularityStats

    stats_dates = select(DemandPopularityStats.biz_dt).distinct()
    stmt = select(func.max(CategoryTreeWeight.biz_dt)).where(
        CategoryTreeWeight.biz_dt.in_(stats_dates)
    )
    return session.scalar(stmt)


def normalize_parent_id(parent_id: int | None) -> int | None:
    if parent_id is None or parent_id == 0:
        return None
    return parent_id


def format_score(avg: Decimal | float | int | None) -> str:
    if avg is None:
        return "—"
    value = float(avg)
    if value >= 100:
        return f"{value:.0f}"
    if value >= 1:
        return f"{value:.1f}"
    return f"{value:.2f}"


def dim_score(
    weight: CategoryTreeWeight | None,
    dim: str,
) -> tuple[float | None, int]:
    """返回 (avg, count)；count>0 即有维度数据（avg 为 0 也算有）。"""
    if weight is None:
        return None, 0
    count = int(getattr(weight, f"{dim}_count", 0) or 0)
    if count <= 0:
        return None, 0
    avg = getattr(weight, f"{dim}_avg", None)
    return (float(avg) if avg is not None else 0.0), count


def build_children_map(
    categories: list[GlobalTreeCategory],
) -> dict[int | None, list[GlobalTreeCategory]]:
    children_map: dict[int | None, list[GlobalTreeCategory]] = {}
    for category in categories:
        parent_key = normalize_parent_id(category.parent_id)
        children_map.setdefault(parent_key, []).append(category)
    for children in children_map.values():
        children.sort(key=lambda c: (c.level or 0, c.id))
    return children_map


def build_score_by_id(
    weights: list[CategoryTreeWeight],
    dim: str,
) -> dict[int, tuple[float | None, int]]:
    return {int(row.category_id): dim_score(row, dim) for row in weights}


def collect_descendant_leaves(
    root_id: int,
    children_map: dict[int | None, list[GlobalTreeCategory]],
) -> list[int]:
    """收集 root 子树内的叶子节点 id（若 root 本身无子节点则含 root）。"""
    leaves: list[int] = []
    stack = [root_id]
    seen: set[int] = set()
    while stack:
        cid = stack.pop()
        if cid in seen:
            continue
        seen.add(cid)
        kids = children_map.get(cid, [])
        if not kids:
            leaves.append(cid)
        else:
            stack.extend(int(c.id) for c in kids)
    leaves.sort()
    return leaves


def collect_descendant_ids(
    root_id: int,
    children_map: dict[int | None, list[GlobalTreeCategory]],
) -> list[int]:
    """收集 root 及其全部子孙节点 id。"""
    ids: list[int] = []
    stack = [root_id]
    seen: set[int] = set()
    while stack:
        cid = stack.pop()
        if cid in seen:
            continue
        seen.add(cid)
        ids.append(cid)
        stack.extend(int(c.id) for c in children_map.get(cid, []))
    ids.sort()
    return ids


def build_category_path(
    category_id: int,
    by_id: dict[int, GlobalTreeCategory],
) -> str | None:
    """从根到指定节点的名称路径，如「美妆护肤 > 护肤 > 防晒」。"""
    cat = by_id.get(category_id)
    if cat is None:
        return None
    names: list[str] = []
    current: GlobalTreeCategory | None = cat
    seen: set[int] = set()
    while current is not None:
        cid = int(current.id)
        if cid in seen:
            break
        seen.add(cid)
        names.append(current.name or str(cid))
        parent_key = normalize_parent_id(current.parent_id)
        current = by_id.get(parent_key) if parent_key is not None else None
    names.reverse()
    return " > ".join(names)