SupplyAgent/supply_infra/scheduler/jobs/expand_demand_from_video_points.py

"""
从 S/A 级需求关联视频中挖掘拓展需求。

任务层负责查库与组装上下文；Agent 仅做语义判断与落库。
"""
from __future__ import annotations

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

from agents.demand_video_expand_agent.run import (
    DemandExpandContext,
    VideoPoint,
    extract_saved_count,
    judge_demand_expansion,
)
from supply_infra.config import get_infra_settings
from supply_infra.db.models.demand_grade import DemandGrade
from supply_infra.db.repositories.demand_grade_repo import DemandGradeRepository
from supply_infra.db.repositories.demand_video_expansion_repo import (
    DemandVideoExpansionRunRepository,
)
from supply_infra.db.repositories.multi_demand_video_point_repo import (
    MultiDemandVideoPointRepository,
)
from supply_infra.db.session import get_session

logger = logging.getLogger(__name__)

_DEFAULT_WORKERS = 5


def _resolve_biz_dt(biz_dt: str | None) -> str:
    if biz_dt:
        text = str(biz_dt).strip()
        if len(text) == 8 and text.isdigit():
            return text
        raise ValueError(f"biz_dt 格式无效，应为 YYYYMMDD: {biz_dt!r}")

    with get_session() as session:
        latest = DemandGradeRepository(session).get_latest_biz_dt()
    if latest:
        return str(latest)

    timezone = ZoneInfo(get_infra_settings().scheduler_timezone)
    return datetime.now(timezone).strftime("%Y%m%d")


def _parse_video_ids(raw: Any) -> list[str]:
    if raw is None:
        return []
    items: list[Any]
    if isinstance(raw, str):
        text = raw.strip()
        if not text:
            return []
        try:
            parsed = json.loads(text)
            items = list(parsed) if isinstance(parsed, list) else [text]
        except (ValueError, TypeError):
            items = [part.strip() for part in text.split(",") if part.strip()]
    elif isinstance(raw, (list, tuple)):
        items = list(raw)
    else:
        return []
    out: list[str] = []
    seen: set[str] = set()
    for item in items:
        vid = str(item).strip() if item is not None else ""
        if not vid or vid in seen:
            continue
        seen.add(vid)
        out.append(vid)
    return out


def _to_float(value: Any) -> float | None:
    if value is None:
        return None
    return float(value)


def load_expand_contexts(
    session,
    biz_dt: str,
    *,
    run_id: str,
    skip_finished: bool = True,
) -> tuple[list[DemandExpandContext], dict[str, int]]:
    """加载待拓展判断的需求上下文，并返回跳过统计。"""
    stats = {
        "total_sa": 0,
        "skipped_no_video": 0,
        "skipped_no_points": 0,
        "skipped_already_done": 0,
    }

    grades = DemandGradeRepository(session).list_by_biz_dt_and_grades(biz_dt, ("S", "A"))
    stats["total_sa"] = len(grades)

    finished_ids: set[int] = set()
    if skip_finished:
        finished_ids = DemandVideoExpansionRunRepository(session).list_finished_grade_ids(
            biz_dt
        )

    rows_with_video: list[tuple[DemandGrade, list[str]]] = []
    all_video_ids: set[str] = set()

    for row in grades:
        if int(row.id) in finished_ids:
            stats["skipped_already_done"] += 1
            continue

        video_ids = _parse_video_ids(row.video_list)
        if not video_ids:
            stats["skipped_no_video"] += 1
            continue

        rows_with_video.append((row, video_ids))
        all_video_ids.update(video_ids)

    points_by_vid = MultiDemandVideoPointRepository(session).list_by_video_ids(all_video_ids)

    contexts: list[DemandExpandContext] = []
    for row, video_ids in rows_with_video:
        points: list[VideoPoint] = []
        for vid in video_ids:
            for point in points_by_vid.get(vid, []):
                points.append(
                    VideoPoint(
                        video_id=vid,
                        point_type=str(point.get("point_type") or ""),
                        point_data=point.get("point_data"),
                        point_desc=point.get("point_desc"),
                    )
                )
        if not points:
            stats["skipped_no_points"] += 1
            continue

        contexts.append(
            DemandExpandContext(
                biz_dt=biz_dt,
                run_id=run_id,
                demand_grade_id=int(row.id),
                demand_name=str(row.demand_name),
                grade=str(row.grade),
                score=_to_float(row.score),
                video_ids=video_ids,
                points=points,
            )
        )

    return contexts, stats


def _record_run(
    *,
    biz_dt: str,
    run_id: str,
    source_demand_grade_id: int,
    saved_count: int,
    status: str,
    error_message: str | None = None,
) -> None:
    with get_session() as session:
        DemandVideoExpansionRunRepository(session).upsert_run(
            biz_dt=biz_dt,
            run_id=run_id,
            source_demand_grade_id=source_demand_grade_id,
            saved_count=saved_count,
            status=status,
            error_message=error_message,
        )


def _process_single_expand(
    ctx: DemandExpandContext,
    *,
    biz_dt: str,
    run_id: str,
) -> dict[str, Any]:
    """并发 worker：对单个需求执行拓展判断并落库执行记录。"""
    try:
        agent_result = judge_demand_expansion(ctx)
        saved_count = extract_saved_count(agent_result)
        _record_run(
            biz_dt=biz_dt,
            run_id=run_id,
            source_demand_grade_id=ctx.demand_grade_id,
            saved_count=saved_count,
            status="finished",
        )
        logger.info(
            "expand demand done: grade_id=%s demand=%s saved=%d iterations=%d",
            ctx.demand_grade_id,
            ctx.demand_name,
            saved_count,
            agent_result.iterations,
        )
        return {
            "success": True,
            "demand_grade_id": ctx.demand_grade_id,
            "demand_name": ctx.demand_name,
            "saved_count": saved_count,
            "iterations": agent_result.iterations,
        }
    except Exception as exc:
        error_text = str(exc)
        _record_run(
            biz_dt=biz_dt,
            run_id=run_id,
            source_demand_grade_id=ctx.demand_grade_id,
            saved_count=0,
            status="failed",
            error_message=error_text,
        )
        logger.exception(
            "expand demand failed: grade_id=%s demand=%s",
            ctx.demand_grade_id,
            ctx.demand_name,
        )
        return {
            "success": False,
            "demand_grade_id": ctx.demand_grade_id,
            "demand_name": ctx.demand_name,
            "error": error_text,
        }


def expand_demand_from_video_points(
    biz_dt: str | None = None,
    *,
    skip_finished: bool = True,
    workers: int = _DEFAULT_WORKERS,
) -> dict[str, Any]:
    """
    对指定业务日的 S/A 需求执行视频点位拓展判断。

    程序负责查需求与点位；无视频或无点位则跳过；有数据则并发调用 Agent。
    """
    started_at = datetime.now()
    run_id = uuid.uuid4().hex

    try:
        resolved_biz_dt = _resolve_biz_dt(biz_dt)
    except Exception as exc:
        logger.exception("expand_demand_from_video_points preflight failed")
        return {
            "success": False,
            "error": str(exc),
            "started_at": started_at.isoformat(),
            "finished_at": datetime.now().isoformat(),
        }

    with get_session() as session:
        contexts, preload_stats = load_expand_contexts(
            session,
            resolved_biz_dt,
            run_id=run_id,
            skip_finished=skip_finished,
        )

    logger.info(
        "expand_demand_from_video_points start: biz_dt=%s run_id=%s workers=%s pending=%s",
        resolved_biz_dt,
        run_id,
        workers,
        len(contexts),
    )

    result: dict[str, Any] = {
        "success": True,
        "run_id": run_id,
        "biz_dt": resolved_biz_dt,
        "started_at": started_at.isoformat(),
        "workers": 0,
        **preload_stats,
        "processed": 0,
        "saved_total": 0,
        "failed": 0,
        "errors": [],
    }

    if not contexts:
        finished_at = datetime.now()
        result["finished_at"] = finished_at.isoformat()
        result["duration_seconds"] = round((finished_at - started_at).total_seconds(), 2)
        logger.info("expand_demand_from_video_points finished: %s", result)
        return result

    worker_count = max(1, min(int(workers), len(contexts)))
    result["workers"] = worker_count

    with ThreadPoolExecutor(max_workers=worker_count) as executor:
        futures = [
            executor.submit(_process_single_expand, ctx, biz_dt=resolved_biz_dt, run_id=run_id)
            for ctx in contexts
        ]
        for future in as_completed(futures):
            try:
                item_result = future.result()
            except Exception as exc:
                logger.exception("expand demand worker 出现未捕获错误: biz_dt=%s", resolved_biz_dt)
                result["failed"] += 1
                result["errors"].append({"error": str(exc)})
                continue

            result["processed"] += 1
            if item_result.get("success"):
                result["saved_total"] += int(item_result.get("saved_count") or 0)
                continue

            result["failed"] += 1
            result["errors"].append(
                {
                    "demand_grade_id": item_result.get("demand_grade_id"),
                    "demand_name": item_result.get("demand_name"),
                    "error": item_result.get("error"),
                }
            )

    finished_at = datetime.now()
    result["finished_at"] = finished_at.isoformat()
    result["duration_seconds"] = round((finished_at - started_at).total_seconds(), 2)
    result["success"] = result["failed"] == 0

    logger.info("expand_demand_from_video_points finished: %s", result)
    return result


if __name__ == "__main__":
    import sys

    from supply_infra.scheduler.cli_result import run_cli

    _args = sys.argv[1:]
    _biz_dt = _args[0] if _args and not _args[0].startswith("-") else None
    _workers_arg = None
    if _biz_dt and len(_args) > 1 and not _args[1].startswith("-"):
        _workers_arg = _args[1]
    _skip_finished = "--force" not in _args

    run_cli(
        lambda: expand_demand_from_video_points(
            _biz_dt,
            skip_finished=_skip_finished,
            workers=int(_workers_arg) if _workers_arg else 5,
        ),
        label="expand_demand_from_video_points",
    )