SupplyAgent/api/services/category_tree.py

"""Build nested category tree from global_tree_category (+ optional weights)."""
from __future__ import annotations

from typing import Any

from supply_infra.db.models.category_tree_weight import CategoryTreeWeight
from supply_infra.db.models.global_tree_category import GlobalTreeCategory
from supply_infra.db.repositories.category_tree_weight_repo import (
    CategoryTreeWeightRepository,
)
from supply_infra.db.repositories.global_tree_category_repo import (
    GlobalTreeCategoryRepository,
)
from supply_infra.db.session import get_session

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

TOTAL_SCORE_KEY = "total_score"
TOTAL_SCORE_DIM_KEYS: tuple[str, ...] = (
    "ext_pop",
    "plat_sust_pop",
    "plat_ly_pop",
    "recent_pop",
)

DIM_META: list[dict[str, str]] = [
    {"key": TOTAL_SCORE_KEY, "label": "全局热度"},
    {"key": "ext_pop", "label": "外部热度"},
    {"key": "plat_sust_pop", "label": "平台持续热度"},
    {"key": "plat_ly_pop", "label": "去年同期热度"},
    {"key": "recent_pop", "label": "近期热度"},
    {"key": "real_rov_7d", "label": "真实ROV(7日)"},
    {"key": "real_vov_7d", "label": "真实VOV(7日)"},
]


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 _weights_payload(row: CategoryTreeWeight | None) -> dict[str, float | None]:
    if row is None:
        return {dim: None for dim in (TOTAL_SCORE_KEY, *DIM_KEYS)}
    return {
        TOTAL_SCORE_KEY: float(row.total_score) if row.total_score is not None else None,
        **{
            dim: (
                float(getattr(row, f"{dim}_avg"))
                if getattr(row, f"{dim}_avg", None) is not None
                else None
            )
            for dim in DIM_KEYS
        },
    }


def _counts_payload(row: CategoryTreeWeight | None) -> dict[str, int]:
    if row is None:
        return {dim: 0 for dim in (TOTAL_SCORE_KEY, *DIM_KEYS)}
    return {
        TOTAL_SCORE_KEY: sum(
            int(getattr(row, f"{dim}_count", 0) or 0) > 0
            for dim in TOTAL_SCORE_DIM_KEYS
        ),
        **{
            dim: int(getattr(row, f"{dim}_count", 0) or 0)
            for dim in DIM_KEYS
        },
    }


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 _to_node(
    category: GlobalTreeCategory,
    children_map: dict[int | None, list[GlobalTreeCategory]],
    weight_by_id: dict[int, CategoryTreeWeight],
) -> dict[str, Any]:
    children = [
        _to_node(child, children_map, weight_by_id)
        for child in children_map.get(category.id, [])
    ]
    weight_row = weight_by_id.get(int(category.id))
    return {
        "id": category.id,
        "name": category.name,
        "level": category.level,
        "description": category.description,
        "weights": _weights_payload(weight_row),
        "counts": _counts_payload(weight_row),
        "hung_word_count": int(weight_row.hung_word_count) if weight_row else 0,
        "children": children,
    }


def build_category_tree(biz_dt: str | None = None) -> dict[str, Any]:
    """Load active categories with per-dim avg; nest as roots.

    Returns ``{biz_dt, dims, nodes}``. When no weight rows exist, ``biz_dt`` is
    ``None`` and each node still carries null ``weights`` / zero ``counts``.
    """
    with get_session() as session:
        categories = GlobalTreeCategoryRepository(session).list_active_categories()
        weight_repo = CategoryTreeWeightRepository(session)
        resolved_dt = biz_dt or weight_repo.get_latest_biz_dt()
        weight_rows: list[CategoryTreeWeight] = []
        if resolved_dt:
            weight_rows = weight_repo.list_by_biz_dt(resolved_dt)
        weight_by_id = {int(r.category_id): r for r in weight_rows}

        children_map = _build_children_map(categories)
        nodes = [
            _to_node(root, children_map, weight_by_id)
            for root in children_map.get(None, [])
        ]

    return {
        "biz_dt": resolved_dt,
        "dims": DIM_META,
        "nodes": nodes,
    }