SupplyAgent/find_agent_v2/demand_context.py

"""Build runnable v2 test contexts directly from generic demand data tables."""

from __future__ import annotations

import json
import uuid
from dataclasses import asdict, dataclass, field
from typing import Any

from sqlalchemy import select

from find_agent_v2.gates import build_rule_snapshot
from find_agent_v2.service import get_find_agent_v2_service
from supply_infra.db.models.demand_grade import DemandGrade
from supply_infra.db.models.demand_video_expansion import DemandVideoExpansion
from supply_infra.db.models.multi_demand_video_detail import MultiDemandVideoDetail
from supply_infra.db.models.multi_demand_video_point import MultiDemandVideoPoint
from supply_infra.db.session import get_session

_POINT_TYPES = {"inspiration", "purpose", "key"}


@dataclass(frozen=True)
class V2ReferencePoint:
    point: str
    point_type: str
    point_desc: str | None = None


@dataclass
class V2ReferenceVideo:
    video_id: str
    title: str
    points: list[V2ReferencePoint] = field(default_factory=list)


@dataclass
class V2DemandContext:
    biz_dt: str
    demand_grade_id: int
    demand_name: str
    grade: str
    score: float | None
    videos: list[V2ReferenceVideo] = field(default_factory=list)

    @property
    def point_count(self) -> int:
        return sum(len(video.points) for video in self.videos)

    def payload(self) -> dict[str, Any]:
        return {
            "biz_dt": self.biz_dt,
            "demand_grade_id": self.demand_grade_id,
            "demand_name": self.demand_name,
            "grade": self.grade,
            "score": self.score,
            "reference_videos": [asdict(video) for video in self.videos],
        }


@dataclass(frozen=True)
class PreparedV2DemandRun:
    run_id: str
    demand_grade_id: int
    demand_name: str
    biz_dt: str
    user_input: str
    reference_video_count: int
    point_count: int

    def summary(self) -> dict[str, Any]:
        return {
            "run_id": self.run_id,
            "demand_grade_id": self.demand_grade_id,
            "demand_name": self.demand_name,
            "biz_dt": self.biz_dt,
            "reference_video_count": self.reference_video_count,
            "point_count": self.point_count,
        }


def _json_ids(raw: str | None) -> list[str]:
    try:
        value = json.loads(raw or "[]")
    except (TypeError, ValueError):
        return []
    if not isinstance(value, list):
        return []
    return list(dict.fromkeys(str(item).strip() for item in value if str(item).strip()))


def _points_from_expansions(rows: list[Any]) -> tuple[list[str], dict[str, list[V2ReferencePoint]]]:
    order: list[str] = []
    result: dict[str, list[V2ReferencePoint]] = {}
    seen: set[tuple[str, str, str]] = set()
    for row in rows:
        video_id = str(row.video_id or "").strip()
        point_type = str(row.point_type or "").strip()
        point = str(row.expanded_text or "").strip()
        key = (video_id, point_type, point)
        if not video_id or point_type not in _POINT_TYPES or not point or key in seen:
            continue
        seen.add(key)
        if video_id not in result:
            order.append(video_id)
            result[video_id] = []
        result[video_id].append(V2ReferencePoint(
            point=point,
            point_type=point_type,
            point_desc=str(row.point_desc or "").strip() or None,
        ))
    return order, result


def _points_from_source_rows(
    video_ids: list[str], rows: list[Any],
) -> tuple[list[str], dict[str, list[V2ReferencePoint]]]:
    result: dict[str, list[V2ReferencePoint]] = {video_id: [] for video_id in video_ids}
    for row in rows:
        video_id = str(row.video_id or "").strip()
        point_type = str(row.point_type or "").strip()
        point = str(row.point_data or "").strip()
        if video_id in result and point_type in _POINT_TYPES and point:
            result[video_id].append(V2ReferencePoint(
                point=point,
                point_type=point_type,
                point_desc=str(row.point_desc or "").strip() or None,
            ))
    return video_ids, {key: value for key, value in result.items() if value}


def _load_context_in_session(session, grade: DemandGrade) -> V2DemandContext | None:
    expansions = list(session.scalars(
        select(DemandVideoExpansion).where(
            DemandVideoExpansion.biz_dt == grade.biz_dt,
            DemandVideoExpansion.source_demand_grade_id == grade.id,
            DemandVideoExpansion.is_delete == 0,
        ).order_by(DemandVideoExpansion.id)
    ))
    video_order, points_by_video = _points_from_expansions(expansions)
    if not points_by_video:
        source_ids = _json_ids(grade.video_list)
        source_points = list(session.scalars(
            select(MultiDemandVideoPoint).where(
                MultiDemandVideoPoint.video_id.in_(source_ids)
            ).order_by(MultiDemandVideoPoint.video_id, MultiDemandVideoPoint.id)
        )) if source_ids else []
        video_order, points_by_video = _points_from_source_rows(source_ids, source_points)
    if not points_by_video:
        return None

    details = {
        str(row.vid): row
        for row in session.scalars(select(MultiDemandVideoDetail).where(
            MultiDemandVideoDetail.vid.in_(list(points_by_video))
        ))
    }
    videos = [
        V2ReferenceVideo(
            video_id=video_id,
            title=str(details[video_id].title or "").strip() or f"(无标题|{video_id})",
            points=points_by_video[video_id],
        )
        for video_id in video_order
        if video_id in points_by_video and video_id in details
    ]
    if not videos:
        return None
    return V2DemandContext(
        biz_dt=str(grade.biz_dt),
        demand_grade_id=int(grade.id),
        demand_name=str(grade.demand_name),
        grade=str(grade.grade),
        score=float(grade.score) if grade.score is not None else None,
        videos=videos,
    )


def load_v2_demand_context(demand_grade_id: int) -> V2DemandContext:
    with get_session() as session:
        grade = session.get(DemandGrade, int(demand_grade_id))
        if grade is None:
            raise LookupError(f"demand_grade_id 不存在: {demand_grade_id}")
        context = _load_context_in_session(session, grade)
        if context is None:
            raise ValueError(f"需求没有可用参考视频和点位: demand_grade_id={demand_grade_id}")
        return context


def pick_latest_v2_demand_context(*, index: int = 0) -> V2DemandContext:
    with get_session() as session:
        latest = session.scalar(select(DemandGrade.biz_dt).where(
            DemandGrade.grade == "S"
        ).order_by(DemandGrade.biz_dt.desc()).limit(1))
        if not latest:
            raise LookupError("本地没有 S 级需求")
        grades = list(session.scalars(select(DemandGrade).where(
            DemandGrade.biz_dt == latest,
            DemandGrade.grade == "S",
        ).order_by(DemandGrade.score.desc(), DemandGrade.id)))
        contexts = [
            context
            for grade in grades
            if (context := _load_context_in_session(session, grade)) is not None
        ]
        if not 0 <= int(index) < len(contexts):
            raise IndexError(f"可用 S 级上下文共 {len(contexts)} 条，index={index} 越界")
        return contexts[int(index)]


def build_v2_user_input(
    context: V2DemandContext, *, run_id: str, rules: dict[str, Any],
) -> str:
    payload = {
        "task": "根据需求语义和参考视频点位，寻找新的高价值抖音候选视频。",
        "run_id": run_id,
        "demand": {
            "biz_dt": context.biz_dt,
            "demand_grade_id": context.demand_grade_id,
            "demand_word": context.demand_name,
            "grade": context.grade,
            "score": context.score,
        },
        "quality_gate_rules": rules,
        "reference_videos": [asdict(video) for video in context.videos],
    }
    return json.dumps(payload, ensure_ascii=False, default=str)


def prepare_v2_demand_run(
    *, demand_grade_id: int | None = None, index: int = 0, run_id: str | None = None,
) -> PreparedV2DemandRun:
    context = (
        load_v2_demand_context(demand_grade_id)
        if demand_grade_id is not None
        else pick_latest_v2_demand_context(index=index)
    )
    run_key = str(run_id or f"local-test-{uuid.uuid4().hex}")[:64]
    rules = build_rule_snapshot()
    user_input = build_v2_user_input(context, run_id=run_key, rules=rules)
    get_find_agent_v2_service().create_run(
        run_id=run_key,
        user_input=user_input,
        demand_word=context.demand_name,
        demand_grade_id=context.demand_grade_id,
        rule_config=rules,
    )
    return PreparedV2DemandRun(
        run_id=run_key,
        demand_grade_id=context.demand_grade_id,
        demand_name=context.demand_name,
        biz_dt=context.biz_dt,
        user_input=user_input,
        reference_video_count=len(context.videos),
        point_count=context.point_count,
    )