SupplyAgent/agents/find_agent/support/douyin_search_tikhub.py

"""通过 TikHub 搜索抖音视频，输出 find_agent 统一候选格式。"""
from __future__ import annotations

import asyncio
import json
import logging
import os
import time
from typing import Any

import httpx
from dotenv import load_dotenv

from agents.find_agent.support.search_persistence import persist_search_payload
from supply_agent.paths import find_project_root

logger = logging.getLogger(__name__)

DOUYIN_SEARCH_TIKHUB_API = (
    "https://api.tikhub.io/api/v1/douyin/search/fetch_video_search_v2"
)
DEFAULT_TIMEOUT = 60.0
DEFAULT_MIN_DURATION_SECONDS = 30
_MIN_REQUEST_INTERVAL_SECONDS = 1.0
_rate_limit_lock = asyncio.Lock()
_last_request_monotonic = 0.0
_env_loaded = False

_CONTENT_TYPE_MAP = {
    "不限": "0",
    "视频": "1",
    "图片": "2",
    "图文": "2",
    "文章": "3",
    "0": "0",
    "1": "1",
    "2": "2",
    "3": "3",
}
_SORT_TYPE_MAP = {
    "综合排序": "0",
    "最多点赞": "1",
    "最新发布": "2",
    "0": "0",
    "1": "1",
    "2": "2",
}
_PUBLISH_TIME_MAP = {
    "不限": "0",
    "一天内": "1",
    "最近一天": "1",
    "一周内": "7",
    "最近一周": "7",
    "半年内": "180",
    "最近半年": "180",
    "0": "0",
    "1": "1",
    "7": "7",
    "180": "180",
}
_DURATION_MAP = {
    "不限": "0",
    "一分钟内": "0-1",
    "1分钟以内": "0-1",
    "1-5分钟": "1-5",
    "五分钟以上": "5-10000",
    "5分钟以上": "5-10000",
    "0": "0",
    "0-1": "0-1",
    "1-5": "1-5",
    "5-10000": "5-10000",
}


def _ensure_env_loaded() -> None:
    global _env_loaded
    if _env_loaded:
        return
    load_dotenv(find_project_root() / ".env")
    _env_loaded = True


def _safe_int(value: Any, default: int = 0) -> int:
    if isinstance(value, bool) or value is None:
        return default
    try:
        return int(float(str(value).strip()))
    except (TypeError, ValueError):
        return default


def _enum_value(value: str, mapping: dict[str, str], field: str) -> str:
    normalized = str(value).strip()
    if normalized not in mapping:
        choices = " / ".join(key for key in mapping if not key.isdigit())
        raise ValueError(f"{field} 不支持「{value}」，可选：{choices}")
    return mapping[normalized]


def _get_aweme_info(item: Any) -> dict[str, Any]:
    if not isinstance(item, dict):
        return {}
    data = item.get("data")
    if not isinstance(data, dict):
        return {}
    aweme_info = data.get("aweme_info")
    return aweme_info if isinstance(aweme_info, dict) else {}


def _extract_topics(aweme: dict[str, Any]) -> list[str]:
    topics: list[str] = []
    for item in aweme.get("topic_list") or []:
        if isinstance(item, str):
            topics.append(item.strip())
        elif isinstance(item, dict):
            topic = (
                item.get("topic_name")
                or item.get("cha_name")
                or item.get("hashtag_name")
                or item.get("name")
            )
            if topic:
                topics.append(str(topic).strip())
    for item in aweme.get("text_extra") or []:
        if isinstance(item, dict):
            topic = item.get("hashtag_name")
            if topic:
                topics.append(str(topic).strip())
    for item in aweme.get("cha_list") or []:
        if isinstance(item, dict):
            topic = item.get("cha_name")
            if topic:
                topics.append(str(topic).strip())
    return list(dict.fromkeys(topic for topic in topics if topic))


def _normalize_aweme(aweme: dict[str, Any]) -> dict[str, Any] | None:
    aweme_id = str(aweme.get("aweme_id") or "").strip()
    if not aweme_id:
        return None
    author = aweme.get("author") if isinstance(aweme.get("author"), dict) else {}
    stats = (
        aweme.get("statistics")
        if isinstance(aweme.get("statistics"), dict)
        else {}
    )
    return {
        "aweme_id": aweme_id,
        "desc": str(
            aweme.get("desc") or aweme.get("item_title") or "无标题"
        )[:200],
        "url": f"https://www.douyin.com/video/{aweme_id}",
        "author": {
            "nickname": str(author.get("nickname") or "未知作者"),
            "sec_uid": str(author.get("sec_uid") or ""),
        },
        "statistics": {
            "digg_count": _safe_int(stats.get("digg_count")),
            "comment_count": _safe_int(stats.get("comment_count")),
            "share_count": _safe_int(stats.get("share_count")),
            "collect_count": _safe_int(stats.get("collect_count")),
            "play_count": _safe_int(stats.get("play_count")),
        },
        "duration_ms": _safe_int(aweme.get("duration")),
        "publish_at": (
            aweme.get("create_time")
            or aweme.get("create_timestamp")
            or aweme.get("publish_timestamp")
        ),
        "topics": _extract_topics(aweme),
    }


def _summary(
    keyword: str,
    results: list[dict[str, Any]],
    *,
    filtered_count: int,
    has_more: bool,
    next_cursor: int,
    search_id: str,
) -> str:
    lines = [
        f"TikHub 搜索关键词「{keyword}」",
        (
            f"保留 {len(results)} 条"
            + (f"，过滤短视频 {filtered_count} 条" if filtered_count else "")
            + (
                f"，还有更多（cursor={next_cursor}, search_id={search_id}）"
                if has_more
                else ""
            )
        ),
        "",
    ]
    for index, item in enumerate(results, 1):
        stats = item["statistics"]
        lines.extend(
            [
                f"{index}. {item['desc'][:50]}",
                f"   ID: {item['aweme_id']}",
                f"   链接: {item['url']}",
                (
                    f"   作者: {item['author']['nickname']} | "
                    f"sec_uid: {item['author']['sec_uid']}"
                ),
                (
                    f"   数据: 点赞 {stats['digg_count']:,} | "
                    f"评论 {stats['comment_count']:,} | "
                    f"分享 {stats['share_count']:,} | "
                    f"收藏 {stats['collect_count']:,}"
                ),
                f"   标签: {'、'.join(item['topics']) or '无'}",
                "",
            ]
        )
    return "\n".join(lines).rstrip()


def _error_result(error: str) -> str:
    return json.dumps(
        {"error": error, "title": "TikHub 抖音搜索失败"},
        ensure_ascii=False,
    )


async def _wait_rate_limit() -> None:
    global _last_request_monotonic
    async with _rate_limit_lock:
        elapsed = time.monotonic() - _last_request_monotonic
        if elapsed < _MIN_REQUEST_INTERVAL_SECONDS:
            await asyncio.sleep(_MIN_REQUEST_INTERVAL_SECONDS - elapsed)
        _last_request_monotonic = time.monotonic()


async def _douyin_search_tikhub_raw(
    keyword: str,
    content_type: str = "视频",
    sort_type: str = "综合排序",
    publish_time: str = "不限",
    cursor: int = 0,
    filter_duration: str = "不限",
    search_id: str = "",
    backtrace: str = "",
    min_duration_seconds: int = DEFAULT_MIN_DURATION_SECONDS,
    timeout: float | None = None,
) -> str:
    """
    使用 TikHub 搜索抖音视频，支持多关键词探索和完整分页状态。

    这是 douyin_search 的独立搜索来源。首次搜索 cursor=0、search_id/backtrace 为空；
    翻页时必须把上次返回的 next_cursor、search_id、backtrace 原样传回。

    Args:
        keyword: Agent 自主确定的实际搜索词。
        content_type: 不限 / 视频 / 图片 / 文章，默认视频；也兼容 TikHub 数字代码。
        sort_type: 综合排序 / 最多点赞 / 最新发布；也兼容 0 / 1 / 2。
        publish_time: 不限 / 一天内 / 一周内 / 半年内；也兼容 0 / 1 / 7 / 180。
        cursor: 首次为 0，翻页使用上次返回的 next_cursor。
        filter_duration: 不限 / 一分钟内 / 1-5分钟 / 5分钟以上。
        search_id: 翻页状态，必须使用同一搜索返回值。
        backtrace: 翻页回溯状态，必须使用同一搜索返回值。
        min_duration_seconds: 客户端最短时长过滤，固定不低于 30 秒。
        timeout: 请求超时秒数，默认 60。

    Returns:
        JSON 字符串。search_results 与 douyin_search 格式兼容，并额外包含
        duration_ms、topics、收藏数和播放数；分页字段为 has_more、next_cursor、
        search_id、backtrace。
    """
    keyword_text = str(keyword).strip()
    if not keyword_text:
        return _error_result("keyword 不能为空")

    try:
        content_type_value = _enum_value(content_type, _CONTENT_TYPE_MAP, "content_type")
        sort_type_value = _enum_value(sort_type, _SORT_TYPE_MAP, "sort_type")
        publish_time_value = _enum_value(
            publish_time, _PUBLISH_TIME_MAP, "publish_time"
        )
        duration_value = _enum_value(
            filter_duration, _DURATION_MAP, "filter_duration"
        )
    except ValueError as exc:
        return _error_result(str(exc))

    _ensure_env_loaded()
    api_key = os.getenv("TIKHUB_API_KEY", "").strip()
    if not api_key:
        return _error_result("未设置环境变量 TIKHUB_API_KEY")

    start_time = time.time()
    request_timeout = timeout if timeout is not None else DEFAULT_TIMEOUT
    payload = {
        "keyword": keyword_text,
        "cursor": max(0, int(cursor)),
        "sort_type": sort_type_value,
        "publish_time": publish_time_value,
        "filter_duration": duration_value,
        "content_type": content_type_value,
        "search_id": str(search_id or ""),
        "backtrace": str(backtrace or ""),
    }

    try:
        await _wait_rate_limit()
        async with httpx.AsyncClient(
            timeout=request_timeout,
            trust_env=False,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {api_key}",
            },
        ) as client:
            response = await client.post(DOUYIN_SEARCH_TIKHUB_API, json=payload)
            response.raise_for_status()
            body = response.json()

        data = body.get("data") if isinstance(body.get("data"), dict) else {}
        business_data = (
            data.get("business_data")
            if isinstance(data.get("business_data"), list)
            else []
        )
        config = (
            data.get("business_config")
            if isinstance(data.get("business_config"), dict)
            else {}
        )
        next_page = (
            config.get("next_page")
            if isinstance(config.get("next_page"), dict)
            else {}
        )

        results: list[dict[str, Any]] = []
        seen: set[str] = set()
        filtered_count = 0
        minimum_ms = max(
            DEFAULT_MIN_DURATION_SECONDS,
            int(min_duration_seconds or DEFAULT_MIN_DURATION_SECONDS),
        ) * 1000
        for raw_item in business_data:
            normalized = _normalize_aweme(_get_aweme_info(raw_item))
            if normalized is None or normalized["aweme_id"] in seen:
                continue
            if (
                minimum_ms
                and normalized["duration_ms"]
                and normalized["duration_ms"] < minimum_ms
            ):
                filtered_count += 1
                continue
            seen.add(normalized["aweme_id"])
            results.append(normalized)

        has_more = bool(config.get("has_more") in (1, True, "1"))
        next_cursor = _safe_int(next_page.get("cursor"))
        next_search_id = str(next_page.get("search_id") or search_id or "")
        next_backtrace = str(
            next_page.get("backtrace") or config.get("backtrace") or backtrace or ""
        )
        duration_ms = int((time.time() - start_time) * 1000)
        result = {
            "title": f"TikHub 抖音搜索: {keyword_text}",
            "output": _summary(
                keyword_text,
                results,
                filtered_count=filtered_count,
                has_more=has_more,
                next_cursor=next_cursor,
                search_id=next_search_id,
            ),
            "provider": "tikhub",
            "keyword": keyword_text,
            "request_params": payload,
            "results_count": len(results),
            "filtered_count": filtered_count,
            "has_more": has_more,
            "next_cursor": next_cursor,
            "search_id": next_search_id,
            "backtrace": next_backtrace,
            "search_results": results,
            "duration_ms": duration_ms,
        }
        logger.info(
            "douyin_search_tikhub completed: keyword=%s results=%d has_more=%s duration_ms=%d",
            keyword_text,
            len(results),
            has_more,
            duration_ms,
        )
        return json.dumps(result, ensure_ascii=False)
    except httpx.HTTPStatusError as exc:
        text = exc.response.text[:1000]
        logger.error(
            "douyin_search_tikhub HTTP error: keyword=%s status=%d",
            keyword_text,
            exc.response.status_code,
        )
        return _error_result(f"HTTP {exc.response.status_code}: {text}")
    except httpx.TimeoutException:
        return _error_result(f"请求超时（{request_timeout}秒）")
    except httpx.RequestError as exc:
        return _error_result(f"网络错误: {exc}")
    except Exception as exc:
        logger.error(
            "douyin_search_tikhub unexpected error: keyword=%s error=%s",
            keyword_text,
            exc,
            exc_info=True,
        )
        return _error_result(f"未知错误: {exc}")


async def douyin_search_tikhub(
    run_id: str,
    keyword: str,
    query_reason: str,
    source_type: str,
    source_value: str | None = None,
    parent_search_id: int | None = None,
    page_no: int = 1,
    content_type: str = "视频",
    sort_type: str = "综合排序",
    publish_time: str = "不限",
    cursor: int = 0,
    filter_duration: str = "不限",
    search_id: str = "",
    backtrace: str = "",
    min_duration_seconds: int = DEFAULT_MIN_DURATION_SECONDS,
    timeout: float | None = None,
) -> str:
    """
    使用 TikHub 搜索一页抖音视频，返回候选基础信息、数据库 ID 和分页状态。
    """
    result = await _douyin_search_tikhub_raw(
        keyword=keyword,
        content_type=content_type,
        sort_type=sort_type,
        publish_time=publish_time,
        cursor=cursor,
        filter_duration=filter_duration,
        search_id=search_id,
        backtrace=backtrace,
        min_duration_seconds=min_duration_seconds,
        timeout=timeout,
    )
    return await asyncio.to_thread(
        persist_search_payload,
        result,
        run_id=run_id,
        keyword=keyword,
        query_reason=query_reason,
        source_type=source_type,
        source_value=source_value,
        parent_search_id=parent_search_id,
        cursor=str(cursor),
        page_no=page_no,
        provider="tikhub",
        content_type=content_type,
        sort_type=sort_type,
        publish_time=publish_time,
        provider_state={"search_id": search_id, "backtrace": backtrace},
    )