"""记录 Agent 制定的批次计划并入库。"""
from __future__ import annotations

import json
import logging
from typing import Any

from agents.demand_grade_orchestrator_agent.common.assignment import (
    MAX_DAILY_BATCHES,
    get_assigned_category_ids,
    get_existing_group_count,
    get_required_hanging_category_ids,
    get_unassigned_hanging_category_ids,
)
from agents.demand_grade_orchestrator_agent.common.plan_persist import persist_groups_one_by_one
from agents.demand_grade_orchestrator_agent.common.plan_record import prepare_grade_groups
from supply_agent.tools import tool

logger = logging.getLogger(__name__)


@tool
def save_grade_plan(
    biz_dt: str,
    grouping_strategy: str,
    groups: list[dict[str, Any]],
) -> str:
    """提交你自行划分的批次计划并入库。可多次调用，每轮可提交多条 group。

    每个 group 需包含：
    - category_ids: 本批分类 ID 列表
    - batch_heat_level: S/A/B/C/D/U
    - planning_reason: 本批规划原因
    - shared_traits: 本批节点共同特征
    """
    existing_groups = get_existing_group_count(biz_dt)
    remaining_quota = max(0, MAX_DAILY_BATCHES - existing_groups)
    required = get_required_hanging_category_ids(biz_dt)
    assigned = get_assigned_category_ids(biz_dt)
    unassigned = get_unassigned_hanging_category_ids(biz_dt)

    base_response: dict[str, Any] = {
        "biz_dt": biz_dt,
        "existing_groups": existing_groups,
        "remaining_batch_quota": remaining_quota,
        "unassigned_category_ids": sorted(unassigned),
        "total_hanging_nodes": len(required),
    }

    if remaining_quota <= 0:
        return json.dumps({
            **base_response,
            "ok": True,
            "persisted": False,
            "message": f"当天批次已达上限 {MAX_DAILY_BATCHES}，本批未入库。",
            "groups": [],
        }, ensure_ascii=False)

    if not groups:
        return json.dumps({
            **base_response,
            "ok": True,
            "persisted": False,
            "message": "groups 为空，未入库。",
            "groups": [],
        }, ensure_ascii=False)

    try:
        prepared = prepare_grade_groups(
            biz_dt,
            grouping_strategy,
            groups,
            assigned_category_ids=assigned,
        )
    except Exception:
        logger.exception("准备批次计划失败: biz_dt=%s", biz_dt)
        return json.dumps({
            **base_response,
            "ok": True,
            "persisted": False,
            "message": "准备批次时发生异常，未入库。",
            "groups": [],
        }, ensure_ascii=False)

    if not prepared.get("groups"):
        return json.dumps({
            **base_response,
            "ok": True,
            "persisted": False,
            "message": "过滤后无有效节点可入库。",
            "groups": [],
            "filtered_category_ids": prepared.get("filtered_category_ids", []),
            "filtered_duplicate_category_ids": prepared.get("filtered_duplicate_category_ids", []),
            "failed_prepare_groups": prepared.get("failed_prepare_groups", []),
        }, ensure_ascii=False)

    persist_result = persist_groups_one_by_one(biz_dt, prepared)

    persisted = persist_result["persisted_group_count"] > 0
    failed_count = len(persist_result["failed_groups"])
    if persisted:
        message = f"已入库 {persist_result['persisted_group_count']} 批"
        if failed_count:
            message += f"，失败 {failed_count} 批"
        message += "。"
    else:
        message = "本批无成功入库记录（可能均已分配、额度已满或全部失败）。"
    return json.dumps({
        "ok": True,
        "persisted": persisted,
        "biz_dt": biz_dt,
        "grouping_strategy": prepared.get("grouping_strategy"),
        "groups": persist_result.get("persisted_groups", []),
        "persisted_group_count": persist_result["persisted_group_count"],
        "skipped_quota": persist_result["skipped_quota"],
        "skipped_empty": persist_result["skipped_empty"],
        "failed_groups": persist_result["failed_groups"],
        "failed_prepare_groups": prepared.get("failed_prepare_groups", []),
        "filtered_category_ids": prepared.get("filtered_category_ids", []),
        "filtered_duplicate_category_ids": prepared.get("filtered_duplicate_category_ids", []),
        "covered_category_ids": [
            cid
            for group in persist_result.get("persisted_groups", [])
            for cid in group.get("category_ids", [])
        ],
        "existing_groups": persist_result["existing_groups"],
        "remaining_batch_quota": persist_result["remaining_batch_quota"],
        "unassigned_category_ids": persist_result["unassigned_category_ids"],
        "coverage_complete": persist_result["coverage_complete"],
        "total_hanging_nodes": persist_result["total_hanging_nodes"],
        "message": message,
    }, ensure_ascii=False)
