SupplyAgent/supply_infra/scheduler/jobs/grade_demand_pool.py

"""计划组落库后执行分级任务（默认定时任务路径为代码自动分配，非统筹 Agent）。"""
from __future__ import annotations

import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from typing import Any
from zoneinfo import ZoneInfo

from agents.demand_grade_agent.run import main as grade_demand_words
from supply_infra.config import get_infra_settings
from supply_infra.scheduler.auto_assign_grade_plan import auto_assign_daily_grade_plan
from supply_infra.db.repositories.demand_grade_plan_repo import DemandGradePlanRepository
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
from supply_infra.scheduler.plan_group_batch import MAX_DEMANDS_PER_BATCH, split_even_batches

logger = logging.getLogger(__name__)

_DEFAULT_WORKERS = 5


def _resolve_biz_dt(biz_dt: str | None) -> str:
    if biz_dt:
        return biz_dt
    return datetime.now(ZoneInfo(get_infra_settings().scheduler_timezone)).strftime("%Y%m%d")


def _materialize_pending_group_items(biz_dt: str) -> int:
    """执行前为 pending 计划组物化待分级需求明细。"""
    with get_session() as session:
        graded_names = DemandGradeRepository(session).get_existing_demand_names(biz_dt)
        return DemandGradePlanRepository(session).materialize_pending_groups(
            biz_dt,
            graded_names=graded_names,
        )


def _group_success(item_summary: dict[str, int], processed_batches: int) -> bool:
    if item_summary.get("pending", 0) > 0:
        return False
    total = sum(item_summary.values())
    if total == 0:
        return True
    if processed_batches > 0:
        return True
    return item_summary.get("failed", 0) == 0 and (
        item_summary.get("finished", 0) > 0 or item_summary.get("skipped", 0) > 0
    )


def _run_group(biz_dt: str, group_id: int, *, max_demands_per_batch: int) -> int:
    """领取并执行一个固定任务，从组内需求明细表按批读取。"""
    with get_session() as session:
        group = DemandGradePlanRepository(session).claim_group(biz_dt, group_id)
    if group is None:
        logger.warning("分级任务已被其他 worker 领取或状态已变化: biz_dt=%s group_id=%s", biz_dt, group_id)
        return 0

    gid = int(group["id"])
    processed_batches = 0
    batch_errors: list[str] = []
    try:
        with get_session() as session:
            items = DemandGradePlanRepository(session).list_pending_group_items(gid)

        for batch_items in split_even_batches(items, max_per_batch=max_demands_per_batch):
            demands = [
                {"pool_id": item["pool_id"], "demand_name": item["demand_name"]}
                for item in batch_items
            ]
            item_ids = [int(item["item_id"]) for item in batch_items]
            try:
                grade_demand_words(demands, biz_dt=biz_dt)
                with get_session() as session:
                    DemandGradePlanRepository(session).mark_group_items_status(
                        item_ids,
                        status="finished",
                    )
                processed_batches += 1
            except Exception as exc:
                logger.exception(
                    "分级子批次失败，跳过并继续本组其余需求: biz_dt=%s group=%s count=%s",
                    biz_dt,
                    group["group_key"],
                    len(demands),
                )
                batch_errors.append(str(exc))
                with get_session() as session:
                    DemandGradePlanRepository(session).mark_group_items_status(
                        item_ids,
                        status="failed",
                        error_message=str(exc),
                    )

        with get_session() as session:
            repo = DemandGradePlanRepository(session)
            item_summary = repo.summarize_group_items(gid)
            repo.finish_group(
                gid,
                success=_group_success(item_summary, processed_batches),
                error_message="; ".join(batch_errors) if batch_errors else None,
            )
    except Exception as exc:
        logger.exception(
            "分级计划任务失败: biz_dt=%s group=%s",
            biz_dt,
            group["group_key"],
        )
        try:
            with get_session() as session:
                DemandGradePlanRepository(session).finish_group(
                    gid,
                    success=False,
                    error_message=str(exc),
                )
        except Exception:
            logger.exception(
                "记录分级计划任务失败状态时发生错误: biz_dt=%s group=%s",
                biz_dt,
                group["group_key"],
            )
        return 0
    return processed_batches


def _execute_plan_tasks(
    biz_dt: str,
    *,
    workers: int,
    max_demands_per_batch: int,
) -> dict[str, Any]:
    """并发执行当天全部 pending 计划任务，仅执行一轮。"""
    with get_session() as session:
        group_ids = DemandGradePlanRepository(session).list_pending_group_ids(biz_dt)
    if not group_ids:
        with get_session() as session:
            final_snapshot = DemandGradePlanRepository(session).get_execution_snapshot(biz_dt)
        return {
            "attempted_groups": 0,
            "groups_run": 0,
            "workers": 0,
            "final_snapshot": final_snapshot,
            "execution_complete": final_snapshot["execution_complete"],
        }

    worker_count = max(1, min(int(workers), len(group_ids)))
    groups_run = 0
    with ThreadPoolExecutor(max_workers=worker_count) as executor:
        futures = [
            executor.submit(_run_group, biz_dt, group_id, max_demands_per_batch=max_demands_per_batch)
            for group_id in group_ids
        ]
        for future in as_completed(futures):
            try:
                groups_run += future.result()
            except Exception:
                logger.exception("分级任务 worker 出现未捕获错误: biz_dt=%s", biz_dt)

    with get_session() as session:
        final_snapshot = DemandGradePlanRepository(session).get_execution_snapshot(biz_dt)
    return {
        "attempted_groups": len(group_ids),
        "groups_run": groups_run,
        "workers": worker_count,
        "final_snapshot": final_snapshot,
        "execution_complete": final_snapshot["execution_complete"],
    }


def execute_plan_tasks_until_complete(
    biz_dt: str,
    *,
    workers: int,
    max_demands_per_batch: int = MAX_DEMANDS_PER_BATCH,
    max_rounds: int = 0,
) -> dict[str, Any]:
    """循环执行 pending 计划任务，直到全部完成或达到 max_rounds。"""
    rounds: list[dict[str, Any]] = []
    groups_run = 0
    round_no = 0
    while True:
        if max_rounds > 0 and round_no >= max_rounds:
            break
        round_no += 1
        result = _execute_plan_tasks(
            biz_dt,
            workers=workers,
            max_demands_per_batch=max_demands_per_batch,
        )
        rounds.append(result)
        groups_run += int(result.get("groups_run") or 0)
        if int(result.get("attempted_groups") or 0) == 0:
            break

    final_snapshot = rounds[-1]["final_snapshot"] if rounds else None
    if final_snapshot is None:
        with get_session() as session:
            final_snapshot = DemandGradePlanRepository(session).get_execution_snapshot(biz_dt)

    return {
        "rounds": round_no,
        "groups_run": groups_run,
        "workers": max(1, int(workers)),
        "round_details": rounds,
        "final_snapshot": final_snapshot,
        "execution_complete": bool(final_snapshot["execution_complete"]),
    }


def _grade_demand_pool_impl(
    resolved_biz_dt: str,
    *,
    workers: int,
    max_demands_per_batch: int = MAX_DEMANDS_PER_BATCH,
    with_orchestrate: bool = True,
    max_rounds: int = 0,
) -> dict[str, Any]:
    with get_session() as session:
        total = MultiDemandPoolDiRepository(session).count_distinct_demand_names(resolved_biz_dt)
        graded_before = DemandGradeRepository(session).count_by_biz_dt(resolved_biz_dt)

    if with_orchestrate:
        try:
            auto_assign_daily_grade_plan(
                biz_dt=resolved_biz_dt,
                max_demands_per_group=max_demands_per_batch,
            )
        except Exception:
            logger.exception("自动分配计划组失败，继续处理数据库中已有任务: biz_dt=%s", resolved_biz_dt)

    materialized = _materialize_pending_group_items(resolved_biz_dt)
    logger.info("物化计划组需求明细: biz_dt=%s items=%s", resolved_biz_dt, materialized)

    plan_execution = execute_plan_tasks_until_complete(
        resolved_biz_dt,
        workers=max(1, int(workers)),
        max_demands_per_batch=max(1, min(int(max_demands_per_batch), MAX_DEMANDS_PER_BATCH)),
        max_rounds=max_rounds,
    )

    with get_session() as session:
        graded_after = DemandGradeRepository(session).count_by_biz_dt(resolved_biz_dt)
    final_snapshot = plan_execution["final_snapshot"]
    result = {
        "success": bool(plan_execution["execution_complete"]),
        "biz_dt": resolved_biz_dt,
        "total": total,
        "graded_before": graded_before,
        "graded_after": graded_after,
        "materialized_items": materialized,
        "planned_category_count": len(final_snapshot["assigned_category_ids"]),
        "planned_groups": final_snapshot["planned_groups"],
        "group_status": final_snapshot["group_status"],
        "plan_execution": plan_execution,
        "workers": plan_execution["workers"],
        "groups_run": plan_execution["groups_run"],
        "run_at": datetime.now().isoformat(),
    }
    logger.info("Grade demand pool completed: %s", result)
    return result


def retry_failed_plan_group_items(
    biz_dt: str | None = None,
    *,
    workers: int = _DEFAULT_WORKERS,
    max_demands_per_batch: int = MAX_DEMANDS_PER_BATCH,
    group_ids: list[int] | None = None,
    dry_run: bool = False,
) -> dict[str, Any]:
    """将 demand_grade_plan_group_item 中 failed 记录重置后重新执行分级。"""
    resolved_biz_dt = _resolve_biz_dt(biz_dt)
    try:
        with get_session() as session:
            failed_items = DemandGradePlanRepository(session).list_failed_group_items(
                resolved_biz_dt,
                group_ids=group_ids,
            )

        if not failed_items:
            with get_session() as session:
                snapshot = DemandGradePlanRepository(session).get_execution_snapshot(resolved_biz_dt)
            return {
                "success": True,
                "biz_dt": resolved_biz_dt,
                "dry_run": dry_run,
                "failed_items": 0,
                "reset": {"reset_items": 0, "reset_groups": 0, "group_ids": [], "items": []},
                "execution": None,
                "final_snapshot": snapshot,
                "run_at": datetime.now().isoformat(),
            }

        if dry_run:
            affected_group_ids = sorted({int(item["group_id"]) for item in failed_items})
            return {
                "success": True,
                "biz_dt": resolved_biz_dt,
                "dry_run": True,
                "failed_items": len(failed_items),
                "reset": {
                    "reset_items": len(failed_items),
                    "reset_groups": len(affected_group_ids),
                    "group_ids": affected_group_ids,
                    "items": failed_items,
                },
                "execution": None,
                "run_at": datetime.now().isoformat(),
            }

        with get_session() as session:
            reset_result = DemandGradePlanRepository(session).reset_failed_items_to_pending(
                resolved_biz_dt,
                group_ids=group_ids,
            )

        with get_session() as session:
            graded_before = DemandGradeRepository(session).count_by_biz_dt(resolved_biz_dt)

        plan_execution = execute_plan_tasks_until_complete(
            resolved_biz_dt,
            workers=max(1, int(workers)),
            max_demands_per_batch=max(1, min(int(max_demands_per_batch), MAX_DEMANDS_PER_BATCH)),
        )

        with get_session() as session:
            graded_after = DemandGradeRepository(session).count_by_biz_dt(resolved_biz_dt)
            remaining_failed = DemandGradePlanRepository(session).list_failed_group_items(
                resolved_biz_dt,
                group_ids=group_ids,
            )

        final_snapshot = plan_execution["final_snapshot"]
        return {
            "success": len(remaining_failed) == 0,
            "biz_dt": resolved_biz_dt,
            "dry_run": False,
            "failed_items": len(failed_items),
            "reset": reset_result,
            "graded_before": graded_before,
            "graded_after": graded_after,
            "remaining_failed": len(remaining_failed),
            "execution": plan_execution,
            "final_snapshot": final_snapshot,
            "group_status": final_snapshot["group_status"],
            "run_at": datetime.now().isoformat(),
        }
    except Exception as exc:
        logger.exception(
            "重试 failed plan group items 发生未捕获错误: biz_dt=%s",
            resolved_biz_dt,
        )
        return {
            "success": False,
            "biz_dt": resolved_biz_dt,
            "error": str(exc),
            "run_at": datetime.now().isoformat(),
        }


def grade_demand_pool(
    biz_dt: str | None = None,
    *,
    workers: int = _DEFAULT_WORKERS,
    max_demands_per_batch: int = MAX_DEMANDS_PER_BATCH,
    with_orchestrate: bool = True,
    max_rounds: int = 0,
) -> dict[str, Any]:
    """执行完整分级闭环；任何错误只记录日志并返回，不向上中断定时任务。"""
    resolved_biz_dt = str(biz_dt or "")
    try:
        resolved_biz_dt = _resolve_biz_dt(biz_dt)
        return _grade_demand_pool_impl(
            resolved_biz_dt,
            workers=workers,
            max_demands_per_batch=max_demands_per_batch,
            with_orchestrate=with_orchestrate,
            max_rounds=max_rounds,
        )
    except Exception as exc:
        logger.exception("需求分级任务发生未捕获错误，已阻止异常中断定时任务: biz_dt=%s", resolved_biz_dt)
        return {
            "success": False,
            "biz_dt": resolved_biz_dt,
            "error": str(exc),
            "run_at": datetime.now().isoformat(),
        }


if __name__ == "__main__":
    import sys

    from supply_infra.scheduler.cli_result import run_cli

    _argv = [a for a in sys.argv[1:] if a != "--retry-failed"]
    _retry_failed = "--retry-failed" in sys.argv[1:]
    _biz_dt = _argv[0] if len(_argv) > 0 else None
    _workers = int(_argv[1]) if len(_argv) > 1 else 5

    def _main() -> dict[str, Any]:
        if _retry_failed:
            return retry_failed_plan_group_items(_biz_dt, workers=_workers)
        return grade_demand_pool(
            _biz_dt,
            workers=_workers,
            with_orchestrate=True,
        )

    run_cli(_main, label="grade_demand_pool")