SupplyAgent/agents/find_agent/tools/video_discovery_store.py

"""持久化 find_agent 的搜索轨迹、候选证据和分池结果。"""
from __future__ import annotations

import hashlib
import json
import logging
import uuid
from decimal import Decimal
from typing import Any

from agents.find_agent.tools.decision_support import (
    audit_video_discovery_process,
)
from supply_agent.tools import tool
from supply_infra.services.video_discovery_service import (
    RunNotFoundError,
    format_db_error,
    get_video_discovery_service,
)

logger = logging.getLogger(__name__)

_RUN_STATUSES = {"running", "finished", "failed"}
_FINAL_DECISION_BUCKETS = {"primary", "rejected"}
_SOURCE_TYPES = {
    "demand",
    "seed",
    "point",
    "tag",
    "author",
    "pagination",
    "mixed",
}
_ROOT_SOURCE_TYPES = {"demand", "seed", "point", "mixed"}


def _json(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, default=str)


def _input_error(message: str) -> str:
    return _json({"error": message, "input_error": True})


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


def _clean_text(value: Any, *, max_length: int | None = None) -> str | None:
    if value is None:
        return None
    text = str(value).strip()
    if not text:
        return None
    return text[:max_length] if max_length else text


def _nonnegative_int(value: Any) -> int | None:
    if value is None or value == "":
        return None
    try:
        return max(0, int(value))
    except (TypeError, ValueError):
        return None


def _optional_decimal(value: Any, places: int) -> Decimal | None:
    if value is None or value == "":
        return None
    try:
        number = Decimal(str(value))
    except (ArithmeticError, TypeError, ValueError):
        return None
    quantum = Decimal(1).scaleb(-places)
    return number.quantize(quantum)


def _search_key(values: dict[str, Any]) -> str:
    identity = {
        key: values.get(key)
        for key in (
            "provider",
            "keyword",
            "content_type",
            "sort_type",
            "publish_time",
            "cursor",
        )
    }
    raw = json.dumps(identity, ensure_ascii=False, sort_keys=True)
    return hashlib.sha256(raw.encode("utf-8")).hexdigest()


def _candidate_from_search_result(item: dict[str, Any], keyword: str) -> dict[str, Any] | None:
    aweme_id = _clean_text(item.get("aweme_id") or item.get("content_id"), max_length=64)
    if not aweme_id:
        return None

    author = item.get("author") if isinstance(item.get("author"), dict) else {}
    stats = item.get("statistics") if isinstance(item.get("statistics"), dict) else {}
    topics = item.get("topics") if isinstance(item.get("topics"), list) else []
    return {
        "aweme_id": aweme_id,
        "title": _clean_text(item.get("desc") or item.get("title"), max_length=512),
        "content_link": _clean_text(
            item.get("url") or item.get("content_link"), max_length=1024
        ),
        "author_name": _clean_text(
            author.get("nickname") or item.get("author_name"), max_length=256
        ),
        "author_sec_uid": _clean_text(
            author.get("sec_uid") or item.get("author_sec_uid"), max_length=256
        ),
        "like_count": _nonnegative_int(
            stats.get("digg_count") or item.get("like_count")
        ),
        "comment_count": _nonnegative_int(
            stats.get("comment_count") or item.get("comment_count")
        ),
        "share_count": _nonnegative_int(
            stats.get("share_count") or item.get("share_count")
        ),
        "collect_count": _nonnegative_int(
            stats.get("collect_count") or item.get("collect_count")
        ),
        "play_count": _nonnegative_int(
            stats.get("play_count") or item.get("play_count")
        ),
        "tags_json": _json(topics) if topics else None,
        "_source_keyword": keyword,
    }


@tool
def create_video_discovery_run(
    demand_word: str,
    seed_video_title: str,
    relevant_points: list[dict[str, Any]],
    seed_video_id: str | None = None,
    demand_grade_id: int | None = None,
    intent_summary: str | None = None,
    run_id: str | None = None,
) -> str:
    """
    创建一次可追踪的视频发现运行。

    若用户消息已提供预创建 run_id，必须原样传入 run_id；工具会复用已有记录，
    不会重复创建。

    Args:
        demand_word: 用户给定需求词；它是输入语义，不强制作为实际搜索词。
        seed_video_title: 与需求相关的参考视频标题。
        relevant_points: 参考视频中与需求相关的点位对象列表。
        seed_video_id: 可选参考视频 id。
        demand_grade_id: 可选 demand_grade.id。
        intent_summary: Agent 对真正受欢迎内容的初步解释，可稍后更新。
        run_id: 系统预创建的运行 id；传入已存在记录时直接复用。

    Returns:
        JSON，包含后续存储工具必须使用的 run_id。
    """
    service = get_video_discovery_service()
    cleaned_run_id = _clean_text(run_id, max_length=64)
    if cleaned_run_id:
        try:
            existing = service.lookup_run(cleaned_run_id)
            if existing is not None:
                return _json(
                    {
                        "title": "视频发现运行已存在",
                        "run_id": cleaned_run_id,
                        "status": existing["status"],
                        "pre_created": True,
                        "output": f"run_id={cleaned_run_id}",
                    }
                )
        except Exception as exc:
            logger.error("create_video_discovery_run lookup failed: %s", exc, exc_info=True)
            return _json({"error": format_db_error(exc), "title": "查询视频发现运行失败"})

    demand = _clean_text(demand_word, max_length=256)
    if not demand:
        return _input_error("demand_word 不能为空")

    new_run_id = cleaned_run_id or uuid.uuid4().hex
    values = {
        "run_id": new_run_id,
        "demand_grade_id": demand_grade_id,
        "demand_word": demand,
        "seed_video_id": _clean_text(seed_video_id, max_length=64),
        "seed_video_title": _clean_text(seed_video_title, max_length=512),
        "relevant_points_json": _json(relevant_points or []),
        "intent_summary": _clean_text(intent_summary),
        "status": "running",
    }
    try:
        created = service.create_run(values)
        return _json(
            {
                "title": "视频发现运行已创建",
                "run_id": created["run_id"],
                "status": created["status"],
                "output": f"run_id={created['run_id']}",
            }
        )
    except Exception as exc:
        logger.error("create_video_discovery_run failed: %s", exc, exc_info=True)
        return _json({"error": format_db_error(exc), "title": "创建视频发现运行失败"})


@tool
def record_video_search_page(
    run_id: str,
    keyword: str,
    query_reason: str,
    source_type: str,
    results: list[dict[str, Any]],
    cursor: str = "0",
    page_no: int = 1,
    has_more: bool = False,
    next_cursor: str | None = None,
    source_value: str | None = None,
    parent_search_id: int | None = None,
    provider: str = "internal_keyword",
    provider_state: dict[str, Any] | None = None,
    content_type: str = "视频",
    sort_type: str = "综合排序",
    publish_time: str = "不限",
    error_message: str | None = None,
) -> str:
    """
    保存一次关键词搜索页，并把该页视频幂等并入候选集。

    每次 douyin_search 后调用。关键词可来自需求语义、参考标题、点位、优质视频标签，
    或前一页的 next_cursor；source_type 用于保留扩展来源。

    Args:
        run_id: create_video_discovery_run 返回值。
        keyword: Agent 本次自主确定的实际搜索词。
        query_reason: 该词验证的内容假设。
        source_type: demand / seed / point / tag / author / pagination / mixed。
        results: douyin_search.search_results 数组。
        cursor / page_no / has_more / next_cursor: 本页翻页状态。
        source_value: 触发扩展的点位、标签或父关键词。
        parent_search_id: 标签扩展或翻页对应的父搜索记录。
        provider: internal_keyword / tikhub / internal_blogger 等来源标识。
        provider_state: 来源特有的分页状态，如 TikHub 的 search_id/backtrace。
        content_type / sort_type / publish_time: 原样保存搜索条件。
        error_message: 搜索失败时保存错误；results 可为空。
    """
    run_text = _clean_text(run_id, max_length=64)
    keyword_text = _clean_text(keyword, max_length=256)
    reason_text = _clean_text(query_reason)
    if not run_text or not keyword_text or not reason_text:
        return _input_error("run_id、keyword、query_reason 不能为空")
    if source_type not in _SOURCE_TYPES:
        return _input_error(f"source_type 必须是: {sorted(_SOURCE_TYPES)}")

    normalized_page_no = max(1, int(page_no))
    normalized_source_type = (
        "pagination" if normalized_page_no > 1 else source_type
    )
    normalized_parent_search_id = (
        None
        if normalized_source_type in _ROOT_SOURCE_TYPES
        else parent_search_id
    )

    candidate_rows = [
        row
        for item in results or []
        if isinstance(item, dict)
        if (row := _candidate_from_search_result(dict(item), keyword_text)) is not None
    ]
    search_values: dict[str, Any] = {
        "run_id": run_text,
        "keyword": keyword_text,
        "query_reason": reason_text,
        "source_type": normalized_source_type,
        "source_value": _clean_text(source_value),
        "parent_search_id": normalized_parent_search_id,
        "provider": _clean_text(provider, max_length=32) or "internal_keyword",
        "provider_state_json": _json(provider_state) if provider_state else None,
        "content_type": _clean_text(content_type, max_length=16) or "视频",
        "sort_type": _clean_text(sort_type, max_length=32) or "综合排序",
        "publish_time": _clean_text(publish_time, max_length=32) or "不限",
        "cursor": _clean_text(cursor, max_length=128) or "0",
        "page_no": normalized_page_no,
        "results_count": len(results or []),
        "new_candidate_count": 0,
        "has_more": int(bool(has_more)),
        "next_cursor": _clean_text(next_cursor, max_length=128),
        "result_ids_json": None,
        "status": "failed" if error_message else "success",
        "error_message": _clean_text(error_message),
    }
    search_values["search_key"] = _search_key(search_values)

    try:
        saved = get_video_discovery_service().save_search_page(
            run_text,
            search_values,
            candidate_rows,
        )
        payload = {
            "title": "搜索页已保存",
            **saved,
            "output": (
                f"search_id={saved['search_id']}，本页 {saved['results_count']} 条，"
                f"新增候选 {saved['new_candidate_count']} 条"
            ),
        }
        return _json(payload)
    except RunNotFoundError as exc:
        return _input_error(str(exc))
    except Exception as exc:
        logger.error("record_video_search_page failed: %s", exc, exc_info=True)
        return _json({"error": format_db_error(exc), "title": "保存搜索页失败"})


def _normalize_evaluation(item: dict[str, Any]) -> dict[str, Any]:
    aweme_id = _clean_text(item.get("aweme_id"), max_length=64)
    if not aweme_id:
        raise ValueError("aweme_id 不能为空")

    content_age = item.get("content_age_evidence")
    account_age = item.get("account_age_evidence")
    age_normalization = item.get("age_normalization")
    detail_verified = bool(item.get("detail_verified"))
    content_portrait_attempted = bool(item.get("content_portrait_attempted"))
    account_portrait_attempted = bool(item.get("account_portrait_attempted"))
    age_portraits_normalized = bool(item.get("age_portraits_normalized"))
    decision_bucket = (
        _clean_text(item.get("decision_bucket"), max_length=24)
        or ""
    )
    if decision_bucket not in _FINAL_DECISION_BUCKETS:
        raise ValueError(
            "decision_bucket 必须是 primary 或 rejected"
        )

    mapping = {
        "aweme_id": aweme_id,
        "title": _clean_text(item.get("title"), max_length=512),
        "content_link": _clean_text(item.get("content_link"), max_length=1024),
        "author_name": _clean_text(item.get("author_name"), max_length=256),
        "author_sec_uid": _clean_text(item.get("author_sec_uid"), max_length=256),
        "source_keywords_json": item.get("source_keywords") or [],
        "source_search_ids_json": item.get("source_search_ids") or [],
        "tags_json": item.get("tags") or [],
        "hit_points_json": item.get("hit_points") or [],
        "play_count": _nonnegative_int(item.get("play_count")),
        "like_count": _nonnegative_int(item.get("like_count")),
        "comment_count": _nonnegative_int(item.get("comment_count")),
        "collect_count": _nonnegative_int(item.get("collect_count")),
        "share_count": _nonnegative_int(item.get("share_count")),
        "publish_timestamp": _nonnegative_int(item.get("publish_timestamp")),
        "content_age_evidence_json": (
            _json(content_age) if content_age is not None else None
        ),
        "account_age_evidence_json": (
            _json(account_age) if account_age is not None else None
        ),
        "age_normalization_json": (
            _json(age_normalization) if age_normalization is not None else None
        ),
        "detail_verified": int(detail_verified),
        "content_portrait_attempted": int(content_portrait_attempted),
        "account_portrait_attempted": int(account_portrait_attempted),
        "age_portraits_normalized": int(age_portraits_normalized),
        "expansion_worthy_tags_json": item.get("expansion_worthy_tags") or [],
        "relevance_score": _optional_decimal(item.get("relevance_score"), 6),
        "elder_score": _optional_decimal(item.get("elder_score"), 6),
        "share_score": _optional_decimal(item.get("share_score"), 6),
        "value_score": _optional_decimal(item.get("value_score"), 2),
        "confidence": _clean_text(item.get("confidence"), max_length=16),
        "relevance_reason": _clean_text(item.get("relevance_reason")),
        "elder_reason": _clean_text(item.get("elder_reason")),
        "share_reason": _clean_text(item.get("share_reason")),
        "decision_reason": _clean_text(item.get("decision_reason")),
        "decision_bucket": decision_bucket,
    }
    return mapping


@tool
def batch_save_video_candidate_evaluations(
    run_id: str,
    items: list[dict[str, Any]],
    run_status: str = "running",
    intent_summary: str | None = None,
    stop_reason: str | None = None,
) -> str:
    """
    原样保存 Agent 给出的候选详情、证据、评分和分池。

    本工具不重算 R/E/S/V，不执行画像证据上限，也不根据阈值修改
    decision_bucket。最终分池只接受 primary / rejected。

    Args:
        run_id: 发现运行 id。
        items: 候选数组。每项至少包含 aweme_id，并由 Agent 直接提供
            decision_bucket。其余详情、证据、评分和理由按模型输出原样保存。
        run_status: running / finished / failed。
        intent_summary: 对目标内容的最终解释。
        stop_reason: 完成或失败时的停止依据。
    """
    run_text = _clean_text(run_id, max_length=64)
    if not run_text:
        return _input_error("run_id 不能为空")
    if run_status not in _RUN_STATUSES:
        return _input_error(
            f"run_status 必须是: {sorted(_RUN_STATUSES)}"
        )

    rows: list[dict[str, Any]] = []
    errors: list[str] = []
    for index, item in enumerate(items or []):
        if not isinstance(item, dict):
            errors.append(f"[{index}] 不是对象")
            continue
        try:
            rows.append(_normalize_evaluation(dict(item)))
        except ValueError as exc:
            errors.append(f"[{index}] {exc}")

    try:
        saved = get_video_discovery_service().save_evaluations_and_finish(
            run_text,
            rows,
            status=run_status,
            intent_summary=_clean_text(intent_summary),
            stop_reason=_clean_text(stop_reason),
        )
        payload = {
            "title": "候选评估已保存",
            "run_id": run_text,
            "saved_count": saved["saved_count"],
            "error_count": len(errors),
            "errors": errors,
            "input_error": bool(errors and not rows),
            "audit_relevant_changed": saved["audit_relevant_changed"],
            "status": saved["status"],
            "search_count": saved["search_count"],
            "primary_count": saved["primary_count"],
            "output": (
                f"保存 {saved['saved_count']} 条；主推荐 {saved['primary_count']} 条；"
                f"状态 {saved['status']}"
            ),
        }
        return _json(payload)
    except RunNotFoundError as exc:
        return _input_error(str(exc))
    except Exception as exc:
        logger.error(
            "batch_save_video_candidate_evaluations failed: %s", exc, exc_info=True
        )
        return _json({"error": format_db_error(exc), "title": "保存候选评估失败"})


@tool
def query_video_discovery_state(
    run_id: str,
    include_rejected: bool = True,
    limit: int = 100,
) -> str:
    """
    查询一次运行已经保存的搜索轨迹、主推荐与淘汰候选。

    用于长搜索过程恢复状态、检查是否真的翻页和扩词，也用于最终自动保留判断。
    """
    run_text = _clean_text(run_id, max_length=64)
    if not run_text:
        return _json({"error": "run_id 不能为空"})
    try:
        state = get_video_discovery_service().get_full_state(
            run_text,
            include_rejected=include_rejected,
            limit=limit,
        )
        run = state["run"]
        payload = {
            "title": f"视频发现状态: {run_text}",
            "run": run,
            "searches": state["searches"],
            "candidates": state["candidates"],
            "output": (
                f"搜索页 {run['search_count']}；主推荐 {run['primary_count']}"
            ),
        }
        return _json(payload)
    except RunNotFoundError:
        return _json({"error": f"run_id 不存在: {run_text}"})
    except Exception as exc:
        logger.error("query_video_discovery_state failed: %s", exc, exc_info=True)
        return _json({"error": format_db_error(exc), "title": "查询视频发现状态失败"})


@tool
def audit_video_discovery_run(
    run_id: str,
    intended_status: str = "finished",
) -> str:
    """
    直接从数据库读取一次发现运行的完整状态并执行结束审计。

    相比把 query_video_discovery_state 的大量 searches/candidates 再复制给审计工具，
    本工具只需要 run_id，可避免长参数截断或 malformed function call。数据库可用时
    应优先使用本工具；数据库不可用的降级流程仍使用 audit_video_discovery_process。

    Args:
        run_id: create_video_discovery_run 返回的运行 id。
        intended_status: 准备结束时传 finished。

    Returns:
        JSON，包含 can_finish、critical_violations、warnings、coverage 和持久化状态。
    """
    run_text = _clean_text(run_id, max_length=64)
    if not run_text:
        return _input_error("run_id 不能为空")
    if intended_status not in _RUN_STATUSES:
        return _input_error(
            f"intended_status 必须是: {sorted(_RUN_STATUSES)}"
        )

    try:
        snapshot = get_video_discovery_service().get_audit_snapshot(run_text)
        persisted_run = snapshot["persisted_run"]
        result = _load_json(
            audit_video_discovery_process(
                searches=snapshot["searches"],
                candidates=snapshot["candidates"],
                intended_status=intended_status,
            ),
            {},
        )
        if not result:
            return _json({"error": "审计工具返回了无效结果"})
        result.update(
            {
                "run_id": run_text,
                "persisted_status": persisted_run["status"],
                "persisted_search_count": persisted_run["search_count"],
                "persisted_primary_count": persisted_run["primary_count"],
            }
        )
        if (
            intended_status == "finished"
            and persisted_run["status"] != "finished"
        ):
            violations = list(result.get("critical_violations") or [])
            violations.append("运行状态尚未持久化为 finished")
            result["critical_violations"] = list(dict.fromkeys(violations))
            result["can_finish"] = False
        return _json(result)
    except RunNotFoundError as exc:
        return _input_error(str(exc))
    except Exception as exc:
        logger.error(
            "audit_video_discovery_run failed: %s",
            exc,
            exc_info=True,
        )
        return _json(
            {"error": format_db_error(exc), "title": "数据库运行审计失败"}
        )