SupplyAgent/agents/demand_grade_agent/tools/batch_save_demand_grades.py

"""
批量保存需求分级结果到 demand_grade 表。
"""
from __future__ import annotations

import logging
from decimal import Decimal
from typing import Any, Optional

from agents.demand_grade_agent.tools.shared import (
    VALID_GRADES,
    collect_strategies,
    dump_int_list,
    merge_video_ids,
    normalize_biz_dt,
)
from supply_agent.tools import tool
from supply_infra.db.repositories.demand_grade_category_rel_repo import (
    DemandGradeCategoryRelRepository,
)
from supply_infra.db.repositories.demand_grade_repo import DemandGradeRepository
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 _optional_decimal(value: Any, field: str, idx: int, errors: list[str]) -> Decimal | None:
    if value is None or value == "":
        return None
    try:
        return Decimal(str(value))
    except Exception:
        errors.append(f"第 {idx} 项 {field} 无效: {value!r}")
        return None


def _optional_int_list(value: Any, field: str, idx: int, errors: list[str]) -> list[int]:
    if value is None:
        return []
    if not isinstance(value, list):
        errors.append(f"第 {idx} 项 {field} 必须是数组: {value!r}")
        return []
    out: list[int] = []
    for v in value:
        try:
            out.append(int(v))
        except (TypeError, ValueError):
            errors.append(f"第 {idx} 项 {field} 含无效元素: {v!r}")
            return []
    return out


def _normalize_items(
    items: list[dict[str, Any]],
    *,
    default_biz_dt: str | None,
) -> tuple[list[dict[str, Any]], list[list[int]], list[list[int]], list[str]]:
    """校验并规范化待落库行，返回 (rows, related_pool_id_lists, category_id_lists, errors)。"""
    rows: list[dict[str, Any]] = []
    related_pool_id_lists: list[list[int]] = []
    category_id_lists: list[list[int]] = []
    errors: list[str] = []
    seen_keys: set[tuple[str, str]] = set()

    for idx, item in enumerate(items):
        if not isinstance(item, dict):
            errors.append(f"第 {idx} 项不是对象")
            continue

        demand_name = str(item.get("demand_name") or "").strip()
        if not demand_name:
            errors.append(f"第 {idx} 项缺少 demand_name")
            continue

        grade = str(item.get("grade") or "").strip().upper()
        if grade not in VALID_GRADES:
            errors.append(f"第 {idx} 项 grade 无效（只能是 {'/'.join(VALID_GRADES)}）: {item.get('grade')!r}")
            continue

        reason = str(item.get("reason") or "").strip()
        if not reason:
            errors.append(f"第 {idx} 项缺少 reason")
            continue

        item_biz_dt, err = normalize_biz_dt(item.get("biz_dt"))
        if err:
            errors.append(f"第 {idx} 项 {err}")
            continue
        resolved_biz_dt = item_biz_dt or default_biz_dt
        if not resolved_biz_dt:
            errors.append(f"第 {idx} 项缺少 biz_dt（且未传入默认 biz_dt）")
            continue

        dedupe_key = (resolved_biz_dt, demand_name)
        if dedupe_key in seen_keys:
            errors.append(f"第 {idx} 项在本次请求中重复: biz_dt={resolved_biz_dt}, demand_name={demand_name}")
            continue

        related_pool_ids = _optional_int_list(item.get("related_pool_ids"), "related_pool_ids", idx, errors)
        if not related_pool_ids:
            errors.append(
                f"第 {idx} 项缺少 related_pool_ids（必填，需先用 search_related_pool_demands 找到对应的 "
                f"multi_demand_pool_di.id）"
            )
            continue

        seen_keys.add(dedupe_key)

        score = _optional_decimal(item.get("score"), "score", idx, errors)
        prior_raw = item.get("prior_total_score")
        if prior_raw is None or prior_raw == "" or prior_raw == "—":
            prior_total_score = None
        else:
            parsed_prior = _optional_decimal(prior_raw, "prior_total_score", idx, errors)
            prior_total_score = None if parsed_prior is None or parsed_prior == 0 else parsed_prior
        posterior_rov_avg = _optional_decimal(item.get("posterior_rov_avg"), "posterior_rov_avg", idx, errors)

        posterior_rov_count = 0
        if item.get("posterior_rov_count") is not None:
            try:
                posterior_rov_count = int(item.get("posterior_rov_count"))
            except (TypeError, ValueError):
                errors.append(f"第 {idx} 项 posterior_rov_count 无效: {item.get('posterior_rov_count')!r}")
                continue

        category_ids = _optional_int_list(item.get("category_ids"), "category_ids", idx, errors)

        has_posterior = 1 if posterior_rov_count > 0 else 0

        rows.append(
            {
                "biz_dt": resolved_biz_dt,
                "demand_name": demand_name,
                "category_ids": dump_int_list(category_ids),
                "grade": grade,
                "score": score,
                "prior_total_score": prior_total_score,
                "posterior_rov_avg": posterior_rov_avg,
                "posterior_rov_count": posterior_rov_count,
                "has_posterior": has_posterior,
                "related_pool_ids": dump_int_list(related_pool_ids),
                "reason": reason,
            }
        )
        related_pool_id_lists.append(related_pool_ids)
        category_id_lists.append(category_ids)

    return rows, related_pool_id_lists, category_id_lists, errors


@tool
def batch_save_demand_grades(items: list[dict[str, Any]], biz_dt: Optional[str] = None) -> str:
    """
    批量保存需求分级结果到 demand_grade 表（按 biz_dt+demand_name upsert，可重复调用覆盖修正）。

    video_list（关联视频列表）与 strategies（来源策略列表）会自动从 related_pool_ids 对应的
    multi_demand_pool_di 原始行推导写入，无需手工传入。category_ids 除了写入展示快照字段，
    也会同步写入 demand_grade_category_rel 映射表，供前端按分类高效查询。

    Args:
        items: 待保存列表，每项字段：
            - demand_name (必填): 需求名称
            - grade (必填): S/A/B/C/D 之一
            - reason (必填): 判断依据，需引用具体的先验/后验数值
            - related_pool_ids (必填): 该需求对应的 multi_demand_pool_di.id 列表，需先调用
              search_related_pool_demands 找到；用于关联原始需求，并自动推导 video_list/strategies
            - score (可选): 0-100 数值分，辅助同级排序
            - category_ids (可选): 归属的树节点 id 列表，会写入 demand_grade_category_rel 映射表
            - prior_total_score (可选): 落库时的先验 total_score 快照
            - posterior_rov_avg / posterior_rov_count (可选): 落库时的后验 real_rov_7d 快照；
              count>0 时自动标记为「有后验数据」
            - biz_dt (可选): 覆盖本项使用的业务日，不传则用调用时的 biz_dt 参数
        biz_dt: 本次调用的默认业务日期 YYYYMMDD；items 内每项也可单独指定 biz_dt 覆盖。

    Returns:
        保存结果摘要，包含成功条数与校验失败说明。
    """
    if not items:
        return "items 不能为空"

    default_biz_dt, err = normalize_biz_dt(biz_dt)
    if err:
        return err

    rows, related_pool_id_lists, category_id_lists, errors = _normalize_items(
        items, default_biz_dt=default_biz_dt
    )
    if not rows:
        detail = "；".join(errors) if errors else "无有效数据"
        return f"没有可保存的数据: {detail}"

    try:
        with get_session() as session:
            pool_repo = MultiDemandPoolDiRepository(session)
            all_pool_ids = sorted({pid for ids in related_pool_id_lists for pid in ids})
            pool_rows = pool_repo.get_by_ids(all_pool_ids) if all_pool_ids else []
            pool_by_id = {int(r.id): r for r in pool_rows}

            final_rows: list[dict[str, Any]] = []
            saved_indices: list[int] = []
            for i, (row, pool_ids) in enumerate(zip(rows, related_pool_id_lists)):
                matched = [pool_by_id[pid] for pid in pool_ids if pid in pool_by_id]
                missing = [pid for pid in pool_ids if pid not in pool_by_id]
                if not matched:
                    errors.append(
                        f"demand_name={row['demand_name']!r} 的 related_pool_ids={pool_ids} "
                        f"均未在 multi_demand_pool_di 中找到，跳过该项"
                    )
                    continue
                if missing:
                    errors.append(
                        f"demand_name={row['demand_name']!r} 的 related_pool_ids 中 {missing} 未找到，已忽略"
                    )
                row["video_list"] = merge_video_ids(matched)
                row["strategies"] = collect_strategies(matched)
                final_rows.append(row)
                saved_indices.append(i)

            if not final_rows:
                detail = "；".join(errors) if errors else "无有效数据"
                return f"没有可保存的数据: {detail}"

            grade_repo = DemandGradeRepository(session)
            affected = grade_repo.bulk_upsert(final_rows)

            names_by_biz_dt: dict[str, list[str]] = {}
            for i in saved_indices:
                names_by_biz_dt.setdefault(rows[i]["biz_dt"], []).append(rows[i]["demand_name"])

            id_by_biz_dt_name: dict[tuple[str, str], int] = {}
            for bd, names in names_by_biz_dt.items():
                for name, demand_grade_id in grade_repo.get_ids_by_names(bd, names).items():
                    id_by_biz_dt_name[(bd, name)] = demand_grade_id

            rel_repo = DemandGradeCategoryRelRepository(session)
            for i in saved_indices:
                key = (rows[i]["biz_dt"], rows[i]["demand_name"])
                demand_grade_id = id_by_biz_dt_name.get(key)
                if demand_grade_id is not None:
                    rel_repo.replace_for_demand_grade(demand_grade_id, category_id_lists[i])

        parts = [f"提交 {len(rows)} 条，成功写入/更新 {affected} 条（{len(final_rows)} 条通过校验）"]
        if errors:
            parts.append(f"校验失败/警告 {len(errors)} 条: " + "；".join(errors))

        message = "。".join(parts)
        logger.info("batch_save_demand_grades completed: %s", message)
        return message

    except Exception as e:
        logger.error("batch_save_demand_grades failed: %s", e, exc_info=True)
        return f"批量保存需求分级失败: {e}"