SupplyAgent/api/services/growth_category_tree.py

"""Build the growth heat map from global_category_v2 and its daily weights."""

from __future__ import annotations

from typing import Any

from sqlalchemy import select

from supply_infra.db.models.global_category_content_weight import (
    GlobalCategoryContentWeight,
)
from supply_infra.db.models.global_v2 import GlobalCategoryV2
from supply_infra.db.repositories.global_category_channel_content_rank_repo import (
    GlobalCategoryChannelContentRankRepository,
)
from supply_infra.db.repositories.global_category_content_weight_repo import (
    GlobalCategoryContentWeightRepository,
)
from supply_infra.db.session import get_session

GROWTH_DIM_META: list[dict[str, str]] = [
    {"key": "read_rate", "label": "阅读率"},
    {"key": "avg_read_rate", "label": "平均阅读率"},
    {"key": "like_rate", "label": "点赞率"},
]
GROWTH_DIM_KEYS = frozenset(dim["key"] for dim in GROWTH_DIM_META)


def _normalize_parent_id(parent_id: int | None, category_ids: set[int]) -> int | None:
    if parent_id in (None, 0) or int(parent_id) not in category_ids:
        return None
    return int(parent_id)


def _build_growth_tree(
    categories: list[GlobalCategoryV2],
    weights: list[GlobalCategoryContentWeight],
) -> list[dict[str, Any]]:
    category_ids = {int(category.stable_id) for category in categories}
    children_by_parent: dict[int | None, list[GlobalCategoryV2]] = {}
    for category in categories:
        parent_id = _normalize_parent_id(category.parent_stable_id, category_ids)
        children_by_parent.setdefault(parent_id, []).append(category)
    for children in children_by_parent.values():
        children.sort(key=lambda row: (row.level or 0, int(row.stable_id)))

    weight_by_stable_id = {int(row.stable_id): row for row in weights}

    def to_node(category: GlobalCategoryV2, ancestors: frozenset[int]) -> dict[str, Any]:
        stable_id = int(category.stable_id)
        weight = weight_by_stable_id.get(stable_id)
        count = int(weight.source_element_count or 0) if weight else 0
        next_ancestors = ancestors | {stable_id}
        children = [
            to_node(child, next_ancestors)
            for child in children_by_parent.get(stable_id, [])
            if int(child.stable_id) not in next_ancestors
        ]
        return {
            "id": stable_id,
            "name": category.name,
            "level": category.level,
            "description": category.description,
            "weights": {
                "read_rate": float(weight.read_rate_score) if weight else None,
                "avg_read_rate": float(weight.avg_read_rate_score) if weight else None,
                "like_rate": float(weight.like_rate_score) if weight else None,
            },
            "counts": {
                "read_rate": count,
                "avg_read_rate": count,
                "like_rate": count,
            },
            "account_uid_count": (
                int(getattr(weight, "account_uid_count", 0) or 0) if weight else 0
            ),
            "channel_content_id_count": (
                int(getattr(weight, "channel_content_id_count", 0) or 0) if weight else 0
            ),
            "cal_fans_num_sum": (
                float(getattr(weight, "cal_fans_num_sum", 0.0) or 0.0) if weight else 0.0
            ),
            "children": children,
        }

    return [to_node(root, frozenset()) for root in children_by_parent.get(None, [])]


def build_growth_category_tree(biz_dt: str | None = None) -> dict[str, Any]:
    """Return the V2 category tree with the requested/latest complete daily scores."""
    with get_session() as session:
        categories = list(
            session.scalars(
                select(GlobalCategoryV2).order_by(
                    GlobalCategoryV2.level, GlobalCategoryV2.stable_id
                )
            ).all()
        )
        weight_repo = GlobalCategoryContentWeightRepository(session)
        resolved_dt = biz_dt or weight_repo.get_latest_completed_biz_dt()
        weights = weight_repo.list_completed_models_by_biz_dt(resolved_dt) if resolved_dt else []
        nodes = _build_growth_tree(categories, weights)

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


def list_growth_category_channel_contents(
    *,
    stable_id: int,
    metric: str,
    biz_dt: str | None = None,
) -> dict[str, Any]:
    if metric not in GROWTH_DIM_KEYS:
        raise ValueError(f"Unsupported growth metric: {metric}")

    with get_session() as session:
        weight_repo = GlobalCategoryContentWeightRepository(session)
        resolved_dt = biz_dt or weight_repo.get_latest_completed_biz_dt()
        rows = (
            GlobalCategoryChannelContentRankRepository(session).list_top(
                biz_dt=resolved_dt,
                stable_id=stable_id,
                metric_type=metric,
            )
            if resolved_dt
            else []
        )
        items = [
            {
                "rank_no": int(row.rank_no),
                "channel_content_id": row.channel_content_id,
                "source_element_id": int(row.source_element_id),
                "source_category_stable_id": int(row.source_category_stable_id),
                "contribution": float(row.contribution),
                "metric_value": float(row.metric_value),
                "weighted_score": float(row.weighted_score),
                "weighted_share": float(row.weighted_share),
            }
            for row in rows
        ]

    return {
        "biz_dt": resolved_dt,
        "stable_id": stable_id,
        "metric": metric,
        "total": len(items),
        "items": items,
    }