"""统筹 Agent 逻辑与工具自检。"""
from __future__ import annotations

import json
import sys
from unittest.mock import patch

from agents.demand_grade_orchestrator_agent.common.assignment import (
    MAX_DAILY_BATCHES,
    dedupe_cross_group_category_ids,
    strip_assigned_category_ids,
)
from agents.demand_grade_orchestrator_agent.common.plan_record import prepare_grade_groups
from agents.demand_grade_orchestrator_agent.run import _summarize_agent_saves
from supply_agent.types import Message, Role


def _ok(name: str) -> None:
    print(f"  ✓ {name}")


def test_prepare_grade_groups_permissive() -> None:
    fake_by_id = {
        101: type("C", (), {"id": 101, "name": "A", "parent_id": None, "level": 1})(),
        102: type("C", (), {"id": 102, "name": "B", "parent_id": None, "level": 1})(),
    }
    fake_weights = {
        101: type("W", (), {"category_id": 101, "total_score": 0.9, "hung_word_count": 5})(),
        102: type("W", (), {"category_id": 102, "total_score": 0.8, "hung_word_count": 0})(),
    }
    groups = [
        {"category_ids": [101, 102, 999], "batch_heat_level": "X", "planning_reason": "", "shared_traits": ""},
        {"category_ids": [101], "batch_heat_level": "A", "planning_reason": "dup", "shared_traits": "dup"},
    ]
    with patch(
        "agents.demand_grade_orchestrator_agent.common.plan_record.load_tree_state",
        return_value=(fake_by_id, {None: [101, 102]}, fake_weights),
    ), patch(
        "agents.demand_grade_orchestrator_agent.common.plan_record.global_heat_positions",
        return_value={101: {"rank": 1, "total": 1, "normalized_score": 0.95}},
    ), patch(
        "agents.demand_grade_orchestrator_agent.common.plan_record.has_hung_demand",
        side_effect=lambda w: w is not None and int(w.hung_word_count or 0) > 0,
    ), patch(
        "agents.demand_grade_orchestrator_agent.common.plan_record.path",
        side_effect=lambda cid, _by: f"path-{cid}",
    ), patch(
        "agents.demand_grade_orchestrator_agent.common.plan_record.heat_level",
        return_value="A",
    ):
        prepared = prepare_grade_groups(
            "20260714",
            "策略",
            groups,
            assigned_category_ids=set(),
        )
    assert len(prepared["groups"]) == 1
    assert prepared["groups"][0]["category_ids"] == [101]
    assert prepared["groups"][0]["batch_heat_level"] == "A"
    assert prepared["groups"][0]["planning_reason"]
    _ok("无需求/非法字段不报错，仅过滤后入库")


def test_dedupe_cross_group() -> None:
    plan = {"groups": [{"category_ids": [1, 2]}, {"category_ids": [2, 3]}]}
    removed = dedupe_cross_group_category_ids(plan)
    assert removed == [2] and plan["groups"][1]["category_ids"] == [3]
    _ok("后批重复节点过滤")


def test_summarize_agent_saves() -> None:
    class Result:
        messages = [
            Message(
                role=Role.TOOL,
                name="save_grade_plan",
                content=json.dumps({"ok": True, "persisted": True, "persisted_group_count": 2}),
            ),
        ]

    summary = _summarize_agent_saves(Result())
    assert summary["save_count"] == 1 and summary["persisted_groups"] == 2
    _ok("统计入库结果")


def test_save_grade_plan_db(biz_dt: str) -> None:
    from agents.demand_grade_orchestrator_agent.tools.save_grade_plan import save_grade_plan
    from agents.demand_grade_orchestrator_agent.common.assignment import resolve_planning_state

    state = resolve_planning_state(biz_dt)
    if not state["unassigned_category_ids"] or state["remaining_batch_quota"] <= 0:
        print("  · save_grade_plan 跳过（无待分配或额度已满）")
        return

    cid = state["unassigned_category_ids"][0]
    with patch(
        "agents.demand_grade_orchestrator_agent.tools.save_grade_plan.persist_groups_one_by_one",
        return_value={
            "persisted_group_count": 1,
            "persisted_groups": [{"category_ids": [cid]}],
            "skipped_quota": 0,
            "skipped_empty": 0,
            "failed_groups": [],
            "existing_groups": 1,
            "remaining_batch_quota": MAX_DAILY_BATCHES - 1,
            "unassigned_category_ids": state["unassigned_category_ids"][1:],
            "coverage_complete": False,
            "total_hanging_nodes": state["total_hanging_nodes"],
        },
    ):
        result = json.loads(
            save_grade_plan(
                biz_dt,
                "自检",
                [{"category_ids": [cid, cid, 999999], "batch_heat_level": "Z"}],
            )
        )
    assert result["ok"] is True and result["persisted"] is True
    _ok("save_grade_plan 宽松校验 + 逐批入库路径")


def main() -> None:
    biz_dt = sys.argv[1] if len(sys.argv) > 1 else "20260714"
    print("=== 统筹 Agent 逻辑自检 ===\n[单元测试]")
    test_prepare_grade_groups_permissive()
    test_dedupe_cross_group()
    test_summarize_agent_saves()
    print("\n[DB 集成检测]")
    test_save_grade_plan_db(biz_dt)
    print("\n全部通过。")


if __name__ == "__main__":
    main()
