SupplyAgent/api/services/video_discovery_records.py

"""Authenticated-user records for the find-agent video discovery workflow."""
from __future__ import annotations

import json
from typing import Any

from sqlalchemy import func, or_, select
from sqlalchemy.orm import load_only

from supply_infra.db.models.video_discovery import (
    VideoDiscoveryCandidate,
    VideoDiscoveryRun,
    VideoDiscoverySearch,
)
from supply_infra.db.session import get_session

_MIN_VISIBLE_BIZ_DT = "20260730"


def _json_value(raw: str | None) -> Any:
    if not raw:
        return None
    try:
        return json.loads(raw)
    except (TypeError, ValueError):
        return raw


def _timestamp(value: Any) -> str | None:
    return value.isoformat() if value is not None else None


def _number(value: Any) -> float | None:
    return float(value) if value is not None else None


def _serialize_run(row: VideoDiscoveryRun, counts: dict[str, int]) -> dict[str, Any]:
    return {
        "id": int(row.id),
        "run_id": row.run_id,
        "biz_dt": row.biz_dt,
        "demand_grade_id": (
            int(row.demand_grade_id) if row.demand_grade_id is not None else None
        ),
        "demand_word": row.demand_word,
        "seed_video_id": row.seed_video_id,
        "seed_video_title": row.seed_video_title,
        "relevant_points": _json_value(row.relevant_points_json),
        "intent_summary": row.intent_summary,
        "status": row.status,
        "search_count": counts.get("search", int(row.search_count or 0)),
        "candidate_count": counts.get("candidate", 0),
        "primary_count": counts.get("primary", int(row.primary_count or 0)),
        "rejected_count": counts.get("rejected", 0),
        "pending_count": counts.get("pending_evaluation", 0),
        "stop_reason": row.stop_reason,
        "rule_version": row.rule_version,
        "rule_config": _json_value(row.rule_config_json),
        "create_time": _timestamp(row.create_time),
        "update_time": _timestamp(row.update_time),
    }


def _serialize_search_candidate(row: VideoDiscoveryCandidate) -> dict[str, Any]:
    return {
        "id": int(row.id),
        "aweme_id": row.aweme_id,
        "title": row.title,
        "content_link": row.content_link,
        "author_name": row.author_name,
        "decision_bucket": row.decision_bucket,
        "publish_at": _timestamp(row.publish_at),
        "duration_seconds": _number(row.duration_seconds),
        "play_count": int(row.play_count) if row.play_count is not None else None,
        "like_count": int(row.like_count) if row.like_count is not None else None,
        "comment_count": (
            int(row.comment_count) if row.comment_count is not None else None
        ),
        "collect_count": (
            int(row.collect_count) if row.collect_count is not None else None
        ),
        "share_count": int(row.share_count) if row.share_count is not None else None,
        "relevance_score": _number(row.relevance_score),
        "elder_score": _number(row.elder_score),
        "share_score": _number(row.share_score),
        "value_score": _number(row.value_score),
        "content_50_plus_ratio": _number(row.content_50_plus_ratio),
        "content_portrait_status": row.content_portrait_status,
        "account_50_plus_ratio": _number(row.account_50_plus_ratio),
        "account_portrait_status": row.account_portrait_status,
        "portrait_conflict": bool(row.portrait_conflict),
        "temporal_status": row.temporal_status,
        "gate_status": row.gate_status,
        "reject_reason_code": row.reject_reason_code,
    }


def _serialize_search(
    row: VideoDiscoverySearch,
    candidates: list[VideoDiscoveryCandidate],
) -> dict[str, Any]:
    return {
        "id": int(row.id),
        "run_id": row.run_id,
        "search_key": row.search_key,
        "keyword": row.keyword,
        "query_reason": row.query_reason,
        "source_type": row.source_type,
        "source_value": row.source_value,
        "parent_search_id": (
            int(row.parent_search_id) if row.parent_search_id is not None else None
        ),
        "provider": row.provider,
        "provider_state": _json_value(row.provider_state_json),
        "content_type": row.content_type,
        "sort_type": row.sort_type,
        "publish_time": row.publish_time,
        "cursor": row.cursor,
        "page_no": int(row.page_no),
        "results_count": int(row.results_count or 0),
        "new_candidate_count": int(row.new_candidate_count or 0),
        "has_more": bool(row.has_more),
        "next_cursor": row.next_cursor,
        "result_ids": _json_value(row.result_ids_json),
        "candidates": [_serialize_search_candidate(candidate) for candidate in candidates],
        "status": row.status,
        "error_message": row.error_message,
        "create_time": _timestamp(row.create_time),
        "update_time": _timestamp(row.update_time),
    }


def _serialize_candidate(row: VideoDiscoveryCandidate) -> dict[str, Any]:
    return {
        "id": int(row.id),
        "run_id": row.run_id,
        "search_id": int(row.search_id) if row.search_id is not None else None,
        "aweme_id": row.aweme_id,
        "title": row.title,
        "content_link": row.content_link,
        "author_name": row.author_name,
        "author_sec_uid": row.author_sec_uid,
        "source_keywords": _json_value(row.source_keywords_json),
        "source_search_ids": _json_value(row.source_search_ids_json),
        "tags": _json_value(row.tags_json),
        "publish_at": _timestamp(row.publish_at),
        "duration_seconds": _number(row.duration_seconds),
        "play_count": int(row.play_count) if row.play_count is not None else None,
        "like_count": int(row.like_count) if row.like_count is not None else None,
        "comment_count": (
            int(row.comment_count) if row.comment_count is not None else None
        ),
        "collect_count": (
            int(row.collect_count) if row.collect_count is not None else None
        ),
        "share_count": int(row.share_count) if row.share_count is not None else None,
        "content_age_evidence": _json_value(row.content_age_evidence_json),
        "account_age_evidence": _json_value(row.account_age_evidence_json),
        "age_normalization": _json_value(row.age_normalization_json),
        "content_50_plus_ratio": _number(row.content_50_plus_ratio),
        "content_50_plus_tgi": _number(row.content_50_plus_tgi),
        "content_portrait_status": row.content_portrait_status,
        "account_50_plus_ratio": _number(row.account_50_plus_ratio),
        "account_50_plus_tgi": _number(row.account_50_plus_tgi),
        "account_portrait_status": row.account_portrait_status,
        "portrait_conflict": bool(row.portrait_conflict),
        "temporal_type": row.temporal_type,
        "temporal_status": row.temporal_status,
        "temporal_evidence": _json_value(row.temporal_evidence_json),
        "gate_status": row.gate_status,
        "gate_results": _json_value(row.gate_results_json),
        "reject_reason_code": row.reject_reason_code,
        "rule_version": row.rule_version,
        "relevance_score": _number(row.relevance_score),
        "elder_score": _number(row.elder_score),
        "share_score": _number(row.share_score),
        "value_score": _number(row.value_score),
        "decision_reason": row.decision_reason,
        "decision_bucket": row.decision_bucket,
        "aigc_crawler_plan_id": row.aigc_crawler_plan_id,
        "aigc_produce_plan_id": row.aigc_produce_plan_id,
        "aigc_publish_plan_id": row.aigc_publish_plan_id,
        "aigc_plan_label": row.aigc_plan_label,
        "create_time": _timestamp(row.create_time),
        "update_time": _timestamp(row.update_time),
    }


def _run_counts(session: Any, run_ids: list[str]) -> dict[str, dict[str, int]]:
    counts = {
        run_id: {
            "search": 0,
            "candidate": 0,
            "primary": 0,
            "rejected": 0,
            "pending_evaluation": 0,
        }
        for run_id in run_ids
    }
    if not run_ids:
        return counts

    search_stmt = (
        select(VideoDiscoverySearch.run_id, func.count(VideoDiscoverySearch.id))
        .where(VideoDiscoverySearch.run_id.in_(run_ids))
        .group_by(VideoDiscoverySearch.run_id)
    )
    for run_id, count in session.execute(search_stmt):
        counts[str(run_id)]["search"] = int(count)

    candidate_stmt = (
        select(
            VideoDiscoveryCandidate.run_id,
            VideoDiscoveryCandidate.decision_bucket,
            func.count(VideoDiscoveryCandidate.id),
        )
        .where(VideoDiscoveryCandidate.run_id.in_(run_ids))
        .group_by(
            VideoDiscoveryCandidate.run_id,
            VideoDiscoveryCandidate.decision_bucket,
        )
    )
    for run_id, bucket, count in session.execute(candidate_stmt):
        run_counts = counts[str(run_id)]
        run_counts[str(bucket)] = int(count)
        run_counts["candidate"] = run_counts.get("candidate", 0) + int(count)
    return counts


def list_video_discovery_runs(
    *,
    biz_dt: str | None = None,
    status: str | None = None,
    keyword: str | None = None,
    limit: int = 20,
    offset: int = 0,
) -> dict[str, Any]:
    """List find-agent runs with live search and candidate counts."""
    with get_session() as session:
        conditions = [VideoDiscoveryRun.biz_dt >= _MIN_VISIBLE_BIZ_DT]
        if biz_dt:
            conditions.append(VideoDiscoveryRun.biz_dt == biz_dt)
        if status:
            conditions.append(VideoDiscoveryRun.status == status)
        normalized_keyword = (keyword or "").strip().lower()
        if normalized_keyword:
            pattern = f"%{normalized_keyword}%"
            conditions.append(
                or_(
                    func.lower(VideoDiscoveryRun.demand_word).like(pattern),
                    func.lower(VideoDiscoveryRun.run_id).like(pattern),
                    func.lower(func.coalesce(VideoDiscoveryRun.seed_video_title, "")).like(
                        pattern
                    ),
                )
            )

        total_stmt = select(func.count(VideoDiscoveryRun.id))
        rows_stmt = select(VideoDiscoveryRun)
        if conditions:
            total_stmt = total_stmt.where(*conditions)
            rows_stmt = rows_stmt.where(*conditions)
        rows_stmt = rows_stmt.order_by(
            VideoDiscoveryRun.create_time.desc(),
            VideoDiscoveryRun.id.desc(),
        ).limit(limit).offset(offset)

        total = int(session.scalar(total_stmt) or 0)
        rows = list(session.scalars(rows_stmt).all())
        counts = _run_counts(session, [row.run_id for row in rows])
        return {
            "items": [_serialize_run(row, counts.get(row.run_id, {})) for row in rows],
            "total": total,
            "limit": limit,
            "offset": offset,
        }


def get_video_discovery_run(run_id: str) -> dict[str, Any] | None:
    with get_session() as session:
        row = session.scalar(
            select(VideoDiscoveryRun).where(
                VideoDiscoveryRun.run_id == run_id,
                VideoDiscoveryRun.biz_dt >= _MIN_VISIBLE_BIZ_DT,
            )
        )
        if row is None:
            return None
        counts = _run_counts(session, [run_id])
        return _serialize_run(row, counts.get(run_id, {}))


def list_video_discovery_searches(
    run_id: str,
    *,
    keyword: str | None = None,
    limit: int = 20,
    offset: int = 0,
) -> dict[str, Any] | None:
    with get_session() as session:
        exists = session.scalar(
            select(VideoDiscoveryRun.id).where(
                VideoDiscoveryRun.run_id == run_id,
                VideoDiscoveryRun.biz_dt >= _MIN_VISIBLE_BIZ_DT,
            )
        )
        if exists is None:
            return None

        conditions = [VideoDiscoverySearch.run_id == run_id]
        normalized_keyword = (keyword or "").strip().lower()
        if normalized_keyword:
            pattern = f"%{normalized_keyword}%"
            conditions.append(
                or_(
                    func.lower(VideoDiscoverySearch.keyword).like(pattern),
                    func.lower(VideoDiscoverySearch.query_reason).like(pattern),
                    func.lower(
                        func.coalesce(VideoDiscoverySearch.source_value, "")
                    ).like(pattern),
                )
            )
        total = int(
            session.scalar(
                select(func.count(VideoDiscoverySearch.id)).where(*conditions)
            )
            or 0
        )
        rows = list(
            session.scalars(
                select(VideoDiscoverySearch)
                .where(*conditions)
                .order_by(VideoDiscoverySearch.id)
                .limit(limit)
                .offset(offset)
            ).all()
        )
        candidates_by_search: dict[int, list[VideoDiscoveryCandidate]] = {
            int(row.id): [] for row in rows
        }
        if candidates_by_search:
            candidate_rows = session.scalars(
                select(VideoDiscoveryCandidate)
                .options(
                    load_only(
                        VideoDiscoveryCandidate.id,
                        VideoDiscoveryCandidate.search_id,
                        VideoDiscoveryCandidate.aweme_id,
                        VideoDiscoveryCandidate.title,
                        VideoDiscoveryCandidate.content_link,
                        VideoDiscoveryCandidate.author_name,
                        VideoDiscoveryCandidate.decision_bucket,
                        VideoDiscoveryCandidate.publish_at,
                        VideoDiscoveryCandidate.duration_seconds,
                        VideoDiscoveryCandidate.play_count,
                        VideoDiscoveryCandidate.like_count,
                        VideoDiscoveryCandidate.comment_count,
                        VideoDiscoveryCandidate.collect_count,
                        VideoDiscoveryCandidate.share_count,
                        VideoDiscoveryCandidate.relevance_score,
                        VideoDiscoveryCandidate.elder_score,
                        VideoDiscoveryCandidate.share_score,
                        VideoDiscoveryCandidate.value_score,
                        VideoDiscoveryCandidate.content_50_plus_ratio,
                        VideoDiscoveryCandidate.content_portrait_status,
                        VideoDiscoveryCandidate.account_50_plus_ratio,
                        VideoDiscoveryCandidate.account_portrait_status,
                        VideoDiscoveryCandidate.portrait_conflict,
                        VideoDiscoveryCandidate.temporal_status,
                        VideoDiscoveryCandidate.gate_status,
                        VideoDiscoveryCandidate.reject_reason_code,
                    )
                )
                .where(VideoDiscoveryCandidate.search_id.in_(candidates_by_search))
                .order_by(
                    VideoDiscoveryCandidate.search_id,
                    VideoDiscoveryCandidate.id,
                )
            ).all()
            for candidate in candidate_rows:
                if candidate.search_id is not None:
                    candidates_by_search[int(candidate.search_id)].append(candidate)
        return {
            "items": [
                _serialize_search(row, candidates_by_search.get(int(row.id), []))
                for row in rows
            ],
            "total": total,
            "limit": limit,
            "offset": offset,
        }


def list_video_discovery_candidates(
    run_id: str,
    *,
    bucket: str | None = None,
    keyword: str | None = None,
    limit: int = 20,
    offset: int = 0,
) -> dict[str, Any] | None:
    with get_session() as session:
        exists = session.scalar(
            select(VideoDiscoveryRun.id).where(
                VideoDiscoveryRun.run_id == run_id,
                VideoDiscoveryRun.biz_dt >= _MIN_VISIBLE_BIZ_DT,
            )
        )
        if exists is None:
            return None

        conditions = [VideoDiscoveryCandidate.run_id == run_id]
        if bucket:
            conditions.append(VideoDiscoveryCandidate.decision_bucket == bucket)
        normalized_keyword = (keyword or "").strip().lower()
        if normalized_keyword:
            pattern = f"%{normalized_keyword}%"
            conditions.append(
                or_(
                    func.lower(VideoDiscoveryCandidate.aweme_id).like(pattern),
                    func.lower(
                        func.coalesce(VideoDiscoveryCandidate.title, "")
                    ).like(pattern),
                    func.lower(
                        func.coalesce(VideoDiscoveryCandidate.author_name, "")
                    ).like(pattern),
                )
            )
        total = int(
            session.scalar(
                select(func.count(VideoDiscoveryCandidate.id)).where(*conditions)
            )
            or 0
        )
        rows = list(
            session.scalars(
                select(VideoDiscoveryCandidate)
                .where(*conditions)
                .order_by(
                    VideoDiscoveryCandidate.value_score.desc(),
                    VideoDiscoveryCandidate.id,
                )
                .limit(limit)
                .offset(offset)
            ).all()
        )
        return {
            "items": [_serialize_candidate(row) for row in rows],
            "total": total,
            "limit": limit,
            "offset": offset,
        }