"""
核心取数工具：需求名 → 归属树节点 → 该节点的先验热度与后验 rov_diff/vov_diff。
"""
from __future__ import annotations

import logging
from typing import Optional

from sqlalchemy.orm import Session

from agents.demand_grade_agent.tools.shared import (
    build_category_path,
    format_category_weight_lines,
    normalize_biz_dt,
    normalize_str_list,
)
from supply_agent.tools import tool
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.demand_belong_category_repo import (
    DemandBelongCategoryRepository,
)
from supply_infra.db.repositories.demand_belong_pool_rel_repo import (
    DemandBelongPoolRelRepository,
)
from supply_infra.db.repositories.global_tree_category_repo import GlobalTreeCategoryRepository
from supply_infra.db.repositories.multi_demand_pool_di_repo import MultiDemandPoolDiRepository
from supply_infra.db.session import get_session

logger = logging.getLogger(__name__)


def _format_weight_row(weight: CategoryTreeWeight, path: str | None) -> str:
    return "\n".join(format_category_weight_lines(
        weight,
        category_id=int(weight.category_id),
        name=None,
        path=path,
        biz_dt=weight.biz_dt,
        hung_word_count=int(weight.hung_word_count or 0),
    ))


def _query_one_demand_category_and_weight(
    session: Session,
    demand_name: str,
    normalized_dt: str | None,
    by_id: dict[int, GlobalTreeCategory],
) -> str:
    belong_ids: set[int] = set()

    pool_repo = MultiDemandPoolDiRepository(session)
    if normalized_dt:
        pool_rows = pool_repo.search_rows_by_name_fragment(normalized_dt, demand_name)
        pool_ids = [row["id"] for row in pool_rows if row["demand_name"] == demand_name]
    else:
        pool_ids = []

    if pool_ids:
        rel_map = DemandBelongPoolRelRepository(session).get_belong_ids_by_pool_ids(pool_ids)
        for ids in rel_map.values():
            belong_ids.update(ids)

    belong_repo = DemandBelongCategoryRepository(session)
    if not belong_ids:
        fuzzy = belong_repo.search_by_name_like(demand_name)
        belong_ids.update(int(row.id) for row in fuzzy)

    if not belong_ids:
        return f"「{demand_name}」未找到归属的树节点（demand_belong_pool_rel 与 demand_belong_category 均无匹配）"

    belong_rows = belong_repo.get_by_ids(list(belong_ids))
    category_ids = sorted({row.category_id for row in belong_rows if row.category_id})
    if not category_ids:
        return f"「{demand_name}」匹配到需求归属记录，但均未挂载 category_id"

    weight_repo = CategoryTreeWeightRepository(session)
    weights = weight_repo.get_by_category_ids(category_ids, normalized_dt)
    if not weights:
        dt_note = f"biz_dt={normalized_dt}" if normalized_dt else "任意业务日"
        return (
            f"「{demand_name}」归属树节点 {category_ids}，"
            f"但 category_tree_weight 在 {dt_note} 无数据"
        )

    lines = [f"「{demand_name}」归属 {len(weights)} 个树节点:"]
    for weight in sorted(
        weights,
        key=lambda w: (w.total_score is None, -(float(w.total_score) if w.total_score is not None else 0)),
    ):
        path = build_category_path(int(weight.category_id), by_id)
        lines.append(_format_weight_row(weight, path))
    return "\n".join(lines)


@tool
def query_demand_category_and_weight(
    demand_names: list[str],
    biz_dt: Optional[str] = None,
) -> str:
    """
    需求名 → 归属树节点 → 全局热度 total_score + 后验 rov_diff/vov_diff（real_rov_7d/real_vov_7d 字段）。

    支持批量传入多个需求名，一次调用返回各词的归属与权重；每段结果前会标注原始 demand_name。

    查找顺序（对每个 demand_name）：
    1. 先在 multi_demand_pool_di 中按需求名精确匹配找到池表行 id；
    2. 经 demand_belong_pool_rel 反查这些池表行归属的 demand_belong_category；
    3. 若第 2 步查不到关系，退化为按需求名模糊匹配 demand_belong_category.name；
    4. 用归属到的 category_id 查询 category_tree_weight 的全局热度与后验数据。

    一个需求可能挂在多个树节点上，会全部列出。

    Args:
        demand_names: 需求名称列表（通常来自 multi_demand_pool_di.demand_name），可一次传多个。
        biz_dt: 业务日期 YYYYMMDD，可选。传入时会先在该日期的需求池里精确匹配需求名，
                再经 demand_belong_pool_rel 反查归属；不传时直接对 demand_belong_category.name
                做模糊匹配（跳过第 1/2 步），权重也取每个树节点各自的最新业务日数据。

    Returns:
        每个 demand_name 一段，段首标注 `--- demand_name: xxx ---`，例如：
        --- demand_name: 减脂期加餐 ---
        「减脂期加餐」归属 1 个树节点:
        [category_id=88] 美食 > 减脂饮食 > 加餐
          biz_dt=20260716  全局热度total_score=3.42
          ...
    """
    normalized_dt, err = normalize_biz_dt(biz_dt)
    if err:
        return err

    names, err = normalize_str_list(demand_names, "demand_names")
    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}

            sections: list[str] = []
            for demand_name in names:
                result = _query_one_demand_category_and_weight(
                    session, demand_name, normalized_dt, by_id
                )
                sections.append(f"--- demand_name: {demand_name} ---\n{result}")

        message = "\n\n".join(sections)
        logger.info(
            "query_demand_category_and_weight completed: count=%d biz_dt=%s",
            len(names),
            normalized_dt,
        )
        return message

    except Exception as e:
        logger.error("query_demand_category_and_weight failed: %s", e, exc_info=True)
        return f"查询需求归属与权重失败: {e}"


def main() -> None:
    print(query_demand_category_and_weight(demand_names=["减脂期加餐", "减脂加餐"]))


if __name__ == "__main__":
    main()
