SupplyAgent/agents/demand_grade_agent/tools/batch_save_demand_grades.py

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

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

from agents.demand_grade_agent.tools.demand_priority import build_demand_priority_index
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__)
_MAX_ERROR_DETAILS = 10


def _format_errors(errors: list[str]) -> str:
    if not errors:
        return ""
    if len(errors) <= _MAX_ERROR_DETAILS:
        return "；".join(errors)
    hidden = len(errors) - _MAX_ERROR_DETAILS
    return "；".join(errors[:_MAX_ERROR_DETAILS]) + f"；...另有 {hidden} 条类似错误"


def _coerce_items(raw: Any) -> tuple[list[Any], str | None]:
    """将 Agent 传入的 items 规范为列表，兼容误传 JSON 字符串。"""
    if raw is None:
        return [], "items 不能为空"
    if isinstance(raw, str):
        text = raw.strip()
        if not text:
            return [], "items 不能为空"
        try:
            raw = json.loads(text)
        except json.JSONDecodeError:
            return [], "items 必须是对象数组，不能把未解析的 JSON 字符串直接传入"
    if isinstance(raw, dict):
        return [raw], None
    if not isinstance(raw, list):
        return [], f"items 必须是数组，当前类型: {type(raw).__name__}"
    return raw, None


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 and item.get("pool_id") is not None:
            try:
                related_pool_ids = [int(item["pool_id"])]
            except (TypeError, ValueError):
                errors.append(f"第 {idx} 项 pool_id 无效: {item.get('pool_id')!r}")
                continue
        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)

        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": None,
                "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 列表；若输入里给了
              pool_id，也可直接写 "pool_id": 123 代替 related_pool_ids
            - score: 无需传入；保存时按当日全量需求池自动计算需求自身来源归一分（0-100），
              即各 strategy 内独立排名归一化后，对该需求已有来源取均值
            - category_ids (可选): 归属的树节点 id 列表，会写入 demand_grade_category_rel 映射表
            - prior_total_score (可选): 落库时的先验 total_score 快照
            - posterior_rov_avg / posterior_rov_count (可选): 落库时的后验 rov_diff 快照（real_rov_7d_avg）；
              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

    coerced_items, coerce_err = _coerce_items(items)
    if coerce_err:
        return coerce_err

    rows, related_pool_id_lists, category_id_lists, errors = _normalize_items(
        coerced_items, default_biz_dt=default_biz_dt
    )
    if not rows:
        detail = _format_errors(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}
            priority_by_biz_dt = {
                resolved_dt: build_demand_priority_index(pool_repo.list_by_biz_dt(resolved_dt))
                for resolved_dt in sorted({row["biz_dt"] for row in 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} 未找到，已忽略"
                    )
                priority = priority_by_biz_dt[row["biz_dt"]].get(row["demand_name"])
                source_rank_score = priority.get("source_rank_score") if priority else None
                row["score"] = (
                    Decimal(str(source_rank_score)) if source_rank_score is not None else None
                )
                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 = _format_errors(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)} 条: {_format_errors(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}"