"""
按分类 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,
    collect_descendant_ids,
    format_score,
    resolve_biz_dt,
)
from supply_agent.tools import tool
from supply_infra.db.models.demand_popularity_stats import DemandPopularityStats
from supply_infra.db.repositories.demand_belong_category_repo import (
    DemandBelongCategoryRepository,
)
from supply_infra.db.repositories.demand_popularity_stats_repo import (
    DemandPopularityStatsRepository,
)
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


def _word_dim_score(
    stats: DemandPopularityStats | None, dim: str
) -> tuple[float | None, int]:
    if stats is None:
        return None, 0
    count = int(getattr(stats, f"{dim}_count", 0) or 0)
    if count <= 0:
        return None, 0
    avg = getattr(stats, f"{dim}_avg", None)
    return (float(avg) if avg is not None else None), count


@tool
def query_demand_words_by_category(
    category_ids: list[int],
    dimension: str,
    include_descendants: bool = True,
    min_count: int = 1,
    top_k: int = 30,
    biz_dt: str | None = None,
) -> str:
    """
    按分类 id 与热度维度，查询挂载需求词及该维词级热度。

    demand_name 必须来自本工具返回的词名（对应 demand_belong_category.name）。
    include_descendants=True 时包含子树内所有挂载词；False 时仅查节点自身直接挂载。
    每个入参 category 下最多返回 top_k 个词，按该维 avg 降序。

    Args:
        category_ids: 分类 id 列表（可多个）。
        dimension: 热度维度，只能是其一：
            - ext_pop：外部热度
            - plat_sust_pop：平台持续热度
            - plat_ly_pop：平台去年同期热度
            - recent_pop：近期热度
        include_descendants: 是否包含子树挂载词，默认 True。
        min_count: 该维 count 下限，默认 1。
        top_k: 每个入参 category 最多返回词数，默认 30。
        biz_dt: 业务日 YYYYMMDD；省略则使用 demand_popularity_stats 最新业务日。

    Returns:
        按入参分类分组的挂载词列表，含 score / count / belong_id。
    """
    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:
        min_count_int = max(0, int(min_count))
    except (TypeError, ValueError):
        return f"min_count 无效: {min_count!r}"
    try:
        top_k_int = max(1, int(top_k))
    except (TypeError, ValueError):
        return f"top_k 无效: {top_k!r}"

    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)

            stats_repo = DemandPopularityStatsRepository(session)
            resolved_biz_dt, err = resolve_biz_dt(
                biz_dt,
                get_latest=stats_repo.get_latest_biz_dt,
                has_data=stats_repo.has_biz_dt,
                table_label="demand_popularity_stats 数据",
            )
            if err:
                return err

            belong_repo = DemandBelongCategoryRepository(session)

            # 收集本次需要查询的全部挂载类目
            scope_by_root: dict[int, list[int]] = {}
            all_scope_ids: set[int] = set()
            for cid in ids:
                if include_descendants:
                    scope = collect_descendant_ids(cid, children_map)
                else:
                    scope = [cid]
                scope_by_root[cid] = scope
                all_scope_ids.update(scope)

            belongs = belong_repo.list_by_category_ids(sorted(all_scope_ids))
            belongs_by_cat: dict[int, list] = {}
            for row in belongs:
                if not row.name:
                    continue
                belongs_by_cat.setdefault(int(row.category_id), []).append(row)

            belong_ids = [int(r.id) for rows in belongs_by_cat.values() for r in rows]
            stats_rows = stats_repo.list_by_biz_dt_and_belong_ids(
                resolved_biz_dt, belong_ids
            )
            stats_by_belong = {
                int(s.demand_category_id): s for s in stats_rows
            }

            lines = [
                f"维度={dim}（{DIM_LABEL[dim]}）| biz_dt={resolved_biz_dt}"
                f" | include_descendants={include_descendants}"
                f" | min_count={min_count_int} | top_k={top_k_int}",
                "",
            ]
            total_words = 0

            for cid in ids:
                cat = by_id.get(cid)
                if cat is None:
                    lines.append(f"[{cid}]")
                    lines.append("  （未找到该分类）")
                    lines.append("")
                    continue

                cname = cat.name or ""
                lines.append(f"[{cid}]{cname}")

                candidates: list[tuple[float, int, int, str, int]] = []
                # (avg, count, belong_id, name, hang_category_id)
                for scope_cid in scope_by_root[cid]:
                    for belong in belongs_by_cat.get(scope_cid, []):
                        avg, count = _word_dim_score(
                            stats_by_belong.get(int(belong.id)), dim
                        )
                        if count < min_count_int:
                            continue
                        score = avg if avg is not None else 0.0
                        candidates.append(
                            (
                                score,
                                count,
                                int(belong.id),
                                str(belong.name),
                                scope_cid,
                            )
                        )

                candidates.sort(key=lambda x: (-x[0], -x[1], x[2]))
                selected = candidates[:top_k_int]
                if not selected:
                    lines.append("  （无符合条件的挂载词）")
                    lines.append("")
                    continue

                total_words += len(selected)
                for avg, count, belong_id, name, hang_cid in selected:
                    hang_note = ""
                    if include_descendants and hang_cid != cid:
                        hang_cat = by_id.get(hang_cid)
                        hang_name = hang_cat.name if hang_cat else ""
                        hang_note = f" hang=[{hang_cid}]{hang_name}"
                    lines.append(
                        f"  - {name}（{format_score(avg)}） count={count}"
                        f"  belong_id={belong_id}{hang_note}"
                    )
                if len(candidates) > top_k_int:
                    lines.append(
                        f"  （另有 {len(candidates) - top_k_int} 个词未展示，已按 avg 截断）"
                    )
                lines.append("")

        logger.info(
            "query_demand_words_by_category completed: dim=%s ids=%s words=%d biz_dt=%s",
            dim,
            ids,
            total_words,
            resolved_biz_dt,
        )
        return "\n".join(lines).rstrip()

    except Exception as e:
        logger.error("query_demand_words_by_category failed: %s", e, exc_info=True)
        return f"查询分类挂载词失败: {e}"


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


if __name__ == "__main__":
    main()
