"""
按分类 id + 热度维度，查询子树内有数据的叶子节点。
"""
from __future__ import annotations

import logging
from typing import Any

from agents.generate_demand_agent.tools.dim_constants import (
    DIM_KEYS,
    DIM_LABEL,
    build_children_map,
    build_score_by_id,
    collect_descendant_leaves,
    format_score,
    resolve_biz_dt,
)
from supply_agent.tools import tool
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

logger = logging.getLogger(__name__)


def _normalize_ids(category_ids: list[Any]) -> tuple[list[int], str | None]:
    if not category_ids:
        return [], "category_ids 不能为空"
    out: list[int] = []
    seen: set[int] = set()
    for raw in category_ids:
        try:
            cid = int(raw)
        except (TypeError, ValueError):
            return [], f"category_ids 含无效 id: {raw!r}"
        if cid in seen:
            continue
        seen.add(cid)
        out.append(cid)
    if not out:
        return [], "category_ids 不能为空"
    return out, None


@tool
def query_category_leaves_by_dim(
    dimension: str,
    category_ids: list[int],
    biz_dt: str | None = None,
) -> str:
    """
    按分类 id 与热度维度，查询各分类节点下有维度数据的叶子节点。

    叶子 = global_tree_category 中无子节点的分类；有维度数据 = 对应维度 count>0
    （avg 为 0 也算有数据）。可一次传入多个 category_id，但只能查一种 dimension。
    返回时按入参分类 id 分组，并带回该分类 id。

    Args:
        dimension: 必选热度维度，只能是其一：
            - ext_pop：外部热度
            - plat_sust_pop：平台持续热度
            - plat_ly_pop：平台去年同期热度
            - recent_pop：近期热度
        category_ids: 要查询的分类 id 列表（可多个）。
        biz_dt: 业务日 YYYYMMDD；省略则使用 category_tree_weight 最新业务日。

    Returns:
        按分类分组的叶子列表，例如：
        维度=recent_pop（近期热度）| biz_dt=20260101
        ## category_id=1 表象
        - [19]表演（3.2） count=5
        - [4]内容形态（1.1） count=2
        ## category_id=2 理念
        （无有数据的叶子节点）
    """
    dim = str(dimension or "").strip()
    if dim not in DIM_KEYS:
        allowed = "、".join(f"{k}（{DIM_LABEL[k]}）" for k in DIM_KEYS)
        return f"dimension 无效，只能是：{allowed}"

    ids, err = _normalize_ids(category_ids)
    if err:
        return err

    try:
        with get_session() as session:
            categories = GlobalTreeCategoryRepository(session).list_active_categories()
            by_id = {int(c.id): c for c in categories}
            children_map = build_children_map(categories)

            weight_repo = CategoryTreeWeightRepository(session)
            resolved_biz_dt, err = resolve_biz_dt(
                biz_dt,
                get_latest=weight_repo.get_latest_biz_dt,
                has_data=weight_repo.has_biz_dt,
                table_label="category_tree_weight 数据",
            )
            if err:
                return err

            score_by_id = build_score_by_id(
                weight_repo.list_by_biz_dt(resolved_biz_dt), dim
            )

            lines = [
                f"维度={dim}（{DIM_LABEL[dim]}）| biz_dt={resolved_biz_dt}",
                "",
            ]
            total_leaves = 0
            for cid in ids:
                cat = by_id.get(cid)
                if cat is None:
                    lines.append(f"## category_id={cid}")
                    lines.append("（未找到该分类）")
                    lines.append("")
                    continue

                cname = cat.name or ""
                lines.append(f"## category_id={cid} {cname}")
                leaf_ids = collect_descendant_leaves(cid, children_map)
                matched: list[str] = []
                for leaf_id in leaf_ids:
                    avg, count = score_by_id.get(leaf_id, (None, 0))
                    if count <= 0:
                        continue
                    leaf = by_id.get(leaf_id)
                    leaf_name = leaf.name if leaf else ""
                    matched.append(
                        f"- [{leaf_id}]{leaf_name}（{format_score(avg)}） count={count}"
                    )
                if matched:
                    total_leaves += len(matched)
                    lines.extend(matched)
                else:
                    lines.append("（无有数据的叶子节点）")
                lines.append("")

        logger.info(
            "query_category_leaves_by_dim completed: dim=%s ids=%s leaves=%d biz_dt=%s",
            dim,
            ids,
            total_leaves,
            resolved_biz_dt,
        )
        return "\n".join(lines).rstrip()

    except Exception as e:
        logger.error("query_category_leaves_by_dim failed: %s", e, exc_info=True)
        return f"查询分类叶子节点失败: {e}"


def main() -> None:
    """本地手动验证。"""
    print(query_category_leaves_by_dim(dimension="recent_pop", category_ids=[1, 2]))


if __name__ == "__main__":
    main()
