"""将 Agent 提交的批次转为可入库结构（仅过滤无需求/已分配节点，不报错）。"""
from __future__ import annotations

import logging
from collections import Counter
from typing import Any

from agents.demand_grade_orchestrator_agent.common.tree_state import (
    HEAT_LEVEL_DEFINITION,
    global_heat_positions,
    has_hung_demand,
    heat_level,
    load_tree_state,
    path,
)

logger = logging.getLogger(__name__)

_LEVEL_ORDER = {level: index for index, level in enumerate(("S", "A", "B", "C", "D", "U"))}
_VALID_LEVELS = frozenset(_LEVEL_ORDER)
_DEFAULT_REASON = "Agent 统筹批次"
_DEFAULT_TRAITS = "按全局树热度与结构分批"


def _parse_category_ids(raw: Any) -> list[int]:
    if not isinstance(raw, list):
        return []
    category_ids: list[int] = []
    for value in raw:
        try:
            category_ids.append(int(value))
        except (TypeError, ValueError):
            continue
    return list(dict.fromkeys(category_ids))


def _normalize_batch_heat_level(raw: Any, category_ids: list[int], positions: dict[int, dict]) -> str:
    level = str(raw or "").strip().upper()
    if level in _VALID_LEVELS:
        return level
    for category_id in category_ids:
        computed = heat_level(positions.get(category_id))
        if computed in _VALID_LEVELS:
            return computed
    return "U"


def _node_position_payload(
    category_id: int,
    by_id: dict[int, Any],
    weights: dict[int, Any],
    positions: dict[int, dict[str, float | int]],
) -> dict[str, Any]:
    position = positions.get(category_id)
    return {
        "category_id": category_id,
        "path": path(category_id, by_id),
        "total_score": (
            float(weights[category_id].total_score)
            if weights.get(category_id) is not None and weights[category_id].total_score is not None
            else None
        ),
        "global_rank": float(position["rank"]) if position is not None else None,
        "global_scored_node_count": int(position["total"]) if position is not None else 0,
        "global_rank_score": float(position["normalized_score"]) if position is not None else None,
    }


def _enrich_group(
    group: dict[str, Any],
    *,
    sequence_no: int,
    by_id: dict[int, Any],
    weights: dict[int, Any],
    positions: dict[int, dict[str, float | int]],
    grouping_strategy: str,
) -> dict[str, Any]:
    category_ids = group["category_ids"]
    batch_heat_level = _normalize_batch_heat_level(group.get("batch_heat_level"), category_ids, positions)
    node_positions = [
        _node_position_payload(category_id, by_id, weights, positions)
        for category_id in category_ids
    ]
    rank_scores = [
        float(item["global_rank_score"])
        for item in node_positions
        if item.get("global_rank_score") is not None
    ]
    raw_scores = [
        float(item["total_score"])
        for item in node_positions
        if item.get("total_score") is not None
    ]
    label = str(HEAT_LEVEL_DEFINITION.get(batch_heat_level, HEAT_LEVEL_DEFINITION["U"])["label"])
    planning_reason = str(group.get("planning_reason") or "").strip() or _DEFAULT_REASON
    shared_traits = str(group.get("shared_traits") or "").strip() or grouping_strategy.strip() or _DEFAULT_TRAITS
    return {
        "group_id": f"{batch_heat_level}-batch-{sequence_no:03d}",
        "sequence_no": sequence_no,
        "category_ids": category_ids,
        "batch_heat_level": batch_heat_level,
        "batch_heat_label": label,
        "batch_global_rank_score": sum(rank_scores) / len(rank_scores) if rank_scores else None,
        "batch_total_score_avg": sum(raw_scores) / len(raw_scores) if raw_scores else None,
        "category_global_positions": node_positions,
        "planning_reason": planning_reason,
        "shared_traits": shared_traits,
    }


def prepare_grade_groups(
    biz_dt: str,
    grouping_strategy: str,
    groups: list[dict[str, Any]],
    *,
    assigned_category_ids: set[int],
) -> dict[str, Any]:
    """过滤并补全批次；仅剔除无需求、不存在、已分配、重复节点。"""
    by_id, _children, weights = load_tree_state(biz_dt)
    positions = global_heat_positions(weights)
    seen: set[int] = set()
    filtered_category_ids: list[int] = []
    filtered_duplicates: list[int] = []
    failed_prepare_groups: list[dict[str, Any]] = []
    draft_groups: list[dict[str, Any]] = []

    for index, raw in enumerate(groups or [], start=1):
        try:
            kept_ids: list[int] = []
            for category_id in _parse_category_ids(raw.get("category_ids")):
                if category_id in seen:
                    filtered_duplicates.append(category_id)
                    continue
                if category_id in assigned_category_ids:
                    filtered_category_ids.append(category_id)
                    continue
                if category_id not in by_id:
                    filtered_category_ids.append(category_id)
                    continue
                if not has_hung_demand(weights.get(category_id)):
                    filtered_category_ids.append(category_id)
                    continue
                seen.add(category_id)
                kept_ids.append(category_id)
            if not kept_ids:
                continue
            draft_groups.append({
                "category_ids": kept_ids,
                "batch_heat_level": raw.get("batch_heat_level"),
                "planning_reason": raw.get("planning_reason"),
                "shared_traits": raw.get("shared_traits"),
                "_sequence_no": index,
            })
        except Exception as exc:
            logger.exception("准备第 %s 批失败，已跳过: biz_dt=%s", index, biz_dt)
            failed_prepare_groups.append({"index": index, "error": str(exc)})

    enriched_groups: list[dict[str, Any]] = []
    for group in draft_groups:
        sequence_no = int(group.get("_sequence_no", len(enriched_groups) + 1))
        try:
            enriched_groups.append(
                _enrich_group(
                    group,
                    sequence_no=sequence_no,
                    by_id=by_id,
                    weights=weights,
                    positions=positions,
                    grouping_strategy=grouping_strategy,
                )
            )
        except Exception as exc:
            logger.exception(
                "补全第 %s 批失败，已跳过: biz_dt=%s category_ids=%s",
                sequence_no,
                biz_dt,
                group.get("category_ids"),
            )
            failed_prepare_groups.append({
                "index": sequence_no,
                "category_ids": group.get("category_ids"),
                "error": str(exc),
            })
    level_counts = Counter(group["batch_heat_level"] for group in enriched_groups)
    return {
        "biz_dt": biz_dt,
        "grouping_strategy": grouping_strategy.strip(),
        "heat_level_definition": HEAT_LEVEL_DEFINITION,
        "global_scored_node_count": len(positions),
        "batch_heat_level_counts": dict(
            sorted(level_counts.items(), key=lambda item: _LEVEL_ORDER[item[0]])
        ),
        "groups": enriched_groups,
        "covered_category_ids": sorted(seen),
        "filtered_category_ids": sorted(set(filtered_category_ids)),
        "filtered_duplicate_category_ids": sorted(set(filtered_duplicates)),
        "failed_prepare_groups": failed_prepare_groups,
    }
