SupplyAgent/agents/find_agent/support/portrait.py

"""
热点宝画像数据工具

调用内部爬虫服务获取账号/内容的粉丝画像。
"""
from __future__ import annotations

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

import httpx

from agents.find_agent.support.age_portrait import normalize_age_portrait_pair

logger = logging.getLogger(__name__)

BATCH_MAX_ITEMS = 8

ACCOUNT_FANS_PORTRAIT_API = (
    "http://crawapi.piaoquantv.com/crawler/dou_yin/re_dian_bao/account_fans_portrait"
)
CONTENT_FANS_PORTRAIT_API = (
    "http://crawapi.piaoquantv.com/crawler/dou_yin/re_dian_bao/video_like_portrait"
)
DEFAULT_TIMEOUT = 60.0


def _top_k(items: dict[str, Any], k: int) -> list[tuple[str, Any]]:
    def percent_value(entry: tuple[str, Any]) -> float:
        metrics = entry[1] if isinstance(entry[1], dict) else {}
        return metrics.get("percentage") or 0.0

    return sorted(items.items(), key=percent_value, reverse=True)[:k]


def _format_portrait_summary(
    header_line: str,
    link_line: str,
    portrait: dict[str, Any],
) -> str:
    summary_lines = [header_line, link_line, ""]
    for key, value in portrait.items():
        if not isinstance(value, dict):
            continue
        if key in ("省份", "城市"):
            summary_lines.append(f"【{key} TOP5】分布")
            items = _top_k(value, 5)
        else:
            summary_lines.append(f"【{key}】分布")
            items = value.items()

        for name, metrics in items:
            ratio = metrics.get("percentage")
            tgi = metrics.get("preference")
            summary_lines.append(f"  {name}: {ratio} (偏好度: {tgi})")
        summary_lines.append("")
    return "\n".join(summary_lines)


def _validate_account_id(account_id: str) -> Optional[str]:
    if not account_id or not isinstance(account_id, str):
        return "account_id 参数无效：必须是非空字符串"
    if not account_id.startswith("MS4wLjABAAAA"):
        return (
            f"account_id 格式错误：必须以 MS4wLjABAAAA 开头，"
            f"当前值: {account_id[:min(20, len(account_id))]}..."
        )
    return None


def _validate_content_id(content_id: str) -> Optional[str]:
    if not content_id or not isinstance(content_id, str):
        return "content_id 参数无效：必须是非空字符串"
    if not content_id.isdigit():
        return f"content_id 格式错误：aweme_id 应该是纯数字，当前值: {content_id[:20]}..."
    if len(content_id) < 15 or len(content_id) > 25:
        return f"content_id 长度异常：期望 15-25 位数字，实际 {len(content_id)} 位"
    return None


def _dimension_flags(
    need_province: bool,
    need_city: bool,
    need_city_level: bool,
    need_gender: bool,
    need_age: bool,
    need_phone_brand: bool,
    need_phone_price: bool,
) -> dict[str, bool]:
    return {
        "need_province": need_province,
        "need_city": need_city,
        "need_city_level": need_city_level,
        "need_gender": need_gender,
        "need_age": need_age,
        "need_phone_brand": need_phone_brand,
        "need_phone_price": need_phone_price,
    }


def _parse_portrait_response(
    data: dict[str, Any],
    *,
    header: str,
    link: str,
) -> dict[str, Any]:
    data_block = data.get("data", {}) if isinstance(data.get("data"), dict) else {}
    portrait = data_block.get("data", {}) if isinstance(data_block.get("data"), dict) else {}
    output = _format_portrait_summary(header, link, portrait)
    has_portrait = bool(portrait and any(isinstance(v, dict) and v for v in portrait.values()))
    return {
        "output": output,
        "has_portrait": has_portrait,
        "portrait_data": portrait,
        "raw_data": data,
    }


async def _fetch_account_portrait(
    client: httpx.AsyncClient,
    account_id: str,
    flags: dict[str, bool],
) -> tuple[Optional[str], Optional[dict[str, Any]]]:
    err = _validate_account_id(account_id)
    if err:
        return err, None

    response = await client.post(
        ACCOUNT_FANS_PORTRAIT_API,
        json={"account_id": account_id, **flags},
        headers={"Content-Type": "application/json"},
    )
    response.raise_for_status()
    data = response.json()

    header = f"账号 {account_id} 的粉丝画像"
    link = (
        f"画像链接：https://douhot.douyin.com/creator/detail?"
        f"active_tab=creator_fans_portrait&creator_id={account_id}"
    )
    return None, _parse_portrait_response(data, header=header, link=link)


async def _fetch_content_portrait(
    client: httpx.AsyncClient,
    content_id: str,
    flags: dict[str, bool],
) -> tuple[Optional[str], Optional[dict[str, Any]]]:
    err = _validate_content_id(content_id)
    if err:
        return err, None

    response = await client.post(
        CONTENT_FANS_PORTRAIT_API,
        json={"content_id": content_id, **flags},
        headers={"Content-Type": "application/json"},
    )
    response.raise_for_status()
    data = response.json()

    header = f"内容 {content_id} 的点赞用户画像"
    link = (
        f"画像链接：https://douhot.douyin.com/video/detail?"
        f"active_tab=video_fans&video_id={content_id}"
    )
    return None, _parse_portrait_response(data, header=header, link=link)


def _success_result(payload: dict[str, Any]) -> str:
    return json.dumps(payload, ensure_ascii=False)


def _error_result(
    error: str,
    *,
    title: str = "画像获取失败",
    input_error: bool = False,
) -> str:
    return json.dumps(
        {"error": error, "title": title, "input_error": input_error},
        ensure_ascii=False,
    )


async def get_account_fans_portrait(
    account_id: str,
    need_province: bool = False,
    need_city: bool = False,
    need_city_level: bool = False,
    need_gender: bool = False,
    need_age: bool = True,
    need_phone_brand: bool = False,
    need_phone_price: bool = False,
    timeout: Optional[float] = None,
) -> str:
    """
    获取抖音账号粉丝画像（热点宝数据）

    获取指定账号的粉丝画像数据，包括年龄、性别、地域等多个维度。

    Args:
        account_id: 抖音账号ID（使用 author.sec_uid）
        need_province: 是否获取省份分布，默认 False
        need_city: 是否获取城市分布，默认 False
        need_city_level: 是否获取城市等级分布（一线/新一线/二线等），默认 False
        need_gender: 是否获取性别分布，默认 False
        need_age: 是否获取年龄分布，默认 True
        need_phone_brand: 是否获取手机品牌分布，默认 False
        need_phone_price: 是否获取手机价格分布，默认 False
        timeout: 超时时间（秒），默认 60

    Returns:
        JSON 字符串，包含 output（文本摘要）、has_portrait、portrait_data、raw_data。
        account_id 使用 author.sec_uid；省份数据只显示 TOP5。
    """
    start_time = time.time()
    request_timeout = timeout if timeout is not None else DEFAULT_TIMEOUT
    flags = _dimension_flags(
        need_province,
        need_city,
        need_city_level,
        need_gender,
        need_age,
        need_phone_brand,
        need_phone_price,
    )

    try:
        async with httpx.AsyncClient(timeout=request_timeout) as client:
            err, ok = await _fetch_account_portrait(client, account_id, flags)

        duration_ms = int((time.time() - start_time) * 1000)

        if err:
            logger.error("get_account_fans_portrait failed: account_id=%s error=%s", account_id, err)
            return _error_result(err, title="账号粉丝画像获取失败")

        assert ok is not None
        logger.info(
            "get_account_fans_portrait completed: account_id=%s has_portrait=%s duration_ms=%d",
            account_id,
            ok["has_portrait"],
            duration_ms,
        )
        return _success_result(
            {
                "title": f"账号粉丝画像: {account_id}",
                "output": ok["output"],
                "has_portrait": ok["has_portrait"],
                "portrait_data": ok["portrait_data"],
                "raw_data": ok["raw_data"],
                "duration_ms": duration_ms,
            }
        )

    except httpx.HTTPStatusError as e:
        logger.error(
            "get_account_fans_portrait HTTP error: account_id=%s status=%d",
            account_id,
            e.response.status_code,
        )
        return _error_result(f"HTTP {e.response.status_code}: {e.response.text}", title="账号粉丝画像获取失败")
    except httpx.TimeoutException:
        logger.error("get_account_fans_portrait timeout: account_id=%s timeout=%s", account_id, request_timeout)
        return _error_result(f"请求超时（{request_timeout}秒）", title="账号粉丝画像获取失败")
    except httpx.RequestError as e:
        logger.error("get_account_fans_portrait network error: account_id=%s error=%s", account_id, e)
        return _error_result(f"网络错误: {e}", title="账号粉丝画像获取失败")
    except Exception as e:
        logger.error(
            "get_account_fans_portrait unexpected error: account_id=%s error=%s",
            account_id,
            e,
            exc_info=True,
        )
        return _error_result(f"未知错误: {e}", title="账号粉丝画像获取失败")


async def get_content_fans_portrait(
    content_id: str,
    need_province: bool = False,
    need_city: bool = False,
    need_city_level: bool = False,
    need_gender: bool = False,
    need_age: bool = True,
    need_phone_brand: bool = False,
    need_phone_price: bool = False,
    timeout: Optional[float] = None,
) -> str:
    """
    获取抖音内容点赞用户画像（热点宝数据）

    获取指定视频内容的点赞用户画像数据，包括年龄、性别、地域等多个维度。

    Args:
        content_id: 抖音内容ID（使用 aweme_id）
        need_province: 是否获取省份分布，默认 False
        need_city: 是否获取城市分布，默认 False
        need_city_level: 是否获取城市等级分布（一线/新一线/二线等），默认 False
        need_gender: 是否获取性别分布，默认 False
        need_age: 是否获取年龄分布，默认 True
        need_phone_brand: 是否获取手机品牌分布，默认 False
        need_phone_price: 是否获取手机价格分布，默认 False
        timeout: 超时时间（秒），默认 60

    Returns:
        JSON 字符串，包含 output（文本摘要）、has_portrait、portrait_data、raw_data。
        若 has_portrait 为 False，可用 get_account_fans_portrait 作为兜底。
    """
    start_time = time.time()
    request_timeout = timeout if timeout is not None else DEFAULT_TIMEOUT
    flags = _dimension_flags(
        need_province,
        need_city,
        need_city_level,
        need_gender,
        need_age,
        need_phone_brand,
        need_phone_price,
    )

    try:
        async with httpx.AsyncClient(timeout=request_timeout) as client:
            err, ok = await _fetch_content_portrait(client, content_id, flags)

        duration_ms = int((time.time() - start_time) * 1000)

        if err:
            logger.error("get_content_fans_portrait failed: content_id=%s error=%s", content_id, err)
            return _error_result(err, title="内容点赞用户画像获取失败")

        assert ok is not None
        logger.info(
            "get_content_fans_portrait completed: content_id=%s has_portrait=%s duration_ms=%d",
            content_id,
            ok["has_portrait"],
            duration_ms,
        )
        return _success_result(
            {
                "title": f"内容点赞用户画像: {content_id}",
                "output": ok["output"],
                "has_portrait": ok["has_portrait"],
                "portrait_data": ok["portrait_data"],
                "raw_data": ok["raw_data"],
                "duration_ms": duration_ms,
            }
        )

    except httpx.HTTPStatusError as e:
        logger.error(
            "get_content_fans_portrait HTTP error: content_id=%s status=%d",
            content_id,
            e.response.status_code,
        )
        return _error_result(f"HTTP {e.response.status_code}: {e.response.text}", title="内容点赞用户画像获取失败")
    except httpx.TimeoutException:
        logger.error("get_content_fans_portrait timeout: content_id=%s timeout=%s", content_id, request_timeout)
        return _error_result(f"请求超时（{request_timeout}秒）", title="内容点赞用户画像获取失败")
    except httpx.RequestError as e:
        logger.error("get_content_fans_portrait network error: content_id=%s error=%s", content_id, e)
        return _error_result(f"网络错误: {e}", title="内容点赞用户画像获取失败")
    except Exception as e:
        logger.error(
            "get_content_fans_portrait unexpected error: content_id=%s error=%s",
            content_id,
            e,
            exc_info=True,
        )
        return _error_result(f"未知错误: {e}", title="内容点赞用户画像获取失败")


async def batch_fetch_portraits(
    candidates_json: str,
    fetch_account_portrait: bool = False,
    need_province: bool = False,
    need_city: bool = False,
    need_city_level: bool = False,
    need_gender: bool = False,
    need_age: bool = True,
    need_phone_brand: bool = False,
    need_phone_price: bool = False,
    timeout: Optional[float] = None,
) -> str:
    """
    批量获取多条候选视频的画像

    依次请求内容点赞画像。fetch_account_portrait=true 时同时请求作者粉丝画像；
    否则仅在内容画像缺失且允许兜底时请求作者画像。
    一次调用返回所有条目，便于比较同一候选的两侧年龄证据。

    Args:
        candidates_json: JSON 数组字符串。每项为对象，字段：
            - aweme_id (必填): 视频 id
            - author_sec_uid (可选): 作者 sec_uid，作者画像或兜底时需要
            - try_account_fallback (可选，默认 true): 为 false 时不请求账号画像
        fetch_account_portrait: 是否为每个候选同时获取作者粉丝画像，默认 False。
            老年受众判断建议设为 True；缺少 author_sec_uid 的条目会跳过作者画像。
        need_* / timeout: 与各单条画像工具一致

    Returns:
        JSON 字符串，包含 output（人类可读摘要）和 results（结构化列表）。
        results 与 candidates 顺序一致，每项含 content / account 子对象。
    """
    start_time = time.time()
    request_timeout = timeout if timeout is not None else DEFAULT_TIMEOUT
    raw = (candidates_json or "").strip()

    if not raw:
        return _error_result(
            "candidates_json 为空",
            title="批量画像失败",
            input_error=True,
        )

    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError as e:
        return _error_result(
            f"candidates_json 不是合法 JSON: {e}",
            title="批量画像失败",
            input_error=True,
        )

    if not isinstance(parsed, list):
        return _error_result(
            "candidates_json 必须是 JSON 数组",
            title="批量画像失败",
            input_error=True,
        )

    if len(parsed) > BATCH_MAX_ITEMS:
        return _error_result(
            f"条目数超过上限 {BATCH_MAX_ITEMS}，请分批调用",
            title="批量画像失败",
            input_error=True,
        )

    flags = _dimension_flags(
        need_province,
        need_city,
        need_city_level,
        need_gender,
        need_age,
        need_phone_brand,
        need_phone_price,
    )

    results: list[dict[str, Any]] = []
    output_chunks: list[str] = []

    try:
        async with httpx.AsyncClient(timeout=request_timeout) as client:
            for idx, entry in enumerate(parsed):
                if not isinstance(entry, dict):
                    results.append(
                        {
                            "aweme_id": None,
                            "error": "条目不是对象",
                            "content": None,
                            "account": None,
                        }
                    )
                    output_chunks.append(f"[{idx}] 跳过：条目不是 JSON 对象")
                    continue

                aweme_id = entry.get("aweme_id") or entry.get("content_id")
                author_sec = entry.get("author_sec_uid") or entry.get("account_id")
                try_fallback = entry.get("try_account_fallback", True)
                if isinstance(try_fallback, str):
                    try_fallback = try_fallback.strip().lower() in ("1", "true", "yes")

                if not aweme_id or not isinstance(aweme_id, str):
                    results.append(
                        {
                            "aweme_id": aweme_id,
                            "error": "缺少 aweme_id",
                            "content": None,
                            "account": None,
                        }
                    )
                    output_chunks.append(f"[{idx}] 跳过：缺少 aweme_id")
                    continue

                item_result: dict[str, Any] = {
                    "aweme_id": aweme_id,
                    "author_sec_uid": author_sec if isinstance(author_sec, str) else None,
                    "try_account_fallback": bool(try_fallback),
                    "fetch_account_portrait": fetch_account_portrait,
                    "content": None,
                    "account": None,
                    "error": None,
                }

                try:
                    cerr, cok = await _fetch_content_portrait(client, aweme_id, flags)
                except httpx.HTTPError as e:
                    cerr, cok = str(e), None

                if cerr:
                    item_result["content"] = {
                        "ok": False,
                        "error": cerr,
                        "has_portrait": False,
                        "portrait_data": {},
                    }
                else:
                    assert cok is not None
                    item_result["content"] = {
                        "ok": True,
                        "error": None,
                        "has_portrait": cok["has_portrait"],
                        "portrait_data": cok["portrait_data"],
                        "output": cok["output"],
                    }

                c_block = item_result["content"]
                content_has = bool(c_block and c_block.get("has_portrait"))
                need_account = fetch_account_portrait or (
                    bool(try_fallback) and not content_has
                )

                if need_account:
                    if not author_sec or not isinstance(author_sec, str):
                        item_result["account"] = {
                            "attempted": False,
                            "skipped_reason": "缺少 author_sec_uid，无法获取作者画像",
                            "has_portrait": False,
                            "portrait_data": {},
                        }
                    else:
                        try:
                            aerr, aok = await _fetch_account_portrait(client, author_sec, flags)
                        except httpx.HTTPError as e:
                            aerr, aok = str(e), None

                        if aerr:
                            item_result["account"] = {
                                "attempted": True,
                                "error": aerr,
                                "has_portrait": False,
                                "portrait_data": {},
                            }
                        else:
                            assert aok is not None
                            item_result["account"] = {
                                "attempted": True,
                                "error": None,
                                "has_portrait": aok["has_portrait"],
                                "portrait_data": aok["portrait_data"],
                                "output": aok["output"],
                            }
                else:
                    skip_reason = (
                        "try_account_fallback 为 false"
                        if not try_fallback
                        else "内容侧已有有效画像，且未要求同时获取作者画像"
                    )
                    item_result["account"] = {
                        "attempted": False,
                        "skipped_reason": skip_reason,
                        "has_portrait": False,
                        "portrait_data": {},
                    }

                content_block = item_result["content"] or {}
                account_block = item_result["account"] or {}
                item_result["age_normalization"] = normalize_age_portrait_pair(
                    content_block.get("portrait_data"),
                    account_block.get("portrait_data"),
                )
                results.append(item_result)
                c_part = item_result["content"] or {}
                a_part = item_result["account"] or {}
                output_chunks.append(
                    f"[{idx}] aweme_id={aweme_id} "
                    f"content_has_portrait={c_part.get('has_portrait')} "
                    f"account_attempted={a_part.get('attempted')} "
                    f"account_has_portrait={a_part.get('has_portrait')}"
                )

        duration_ms = int((time.time() - start_time) * 1000)
        logger.info(
            "batch_fetch_portraits completed: count=%d candidates=%d duration_ms=%d",
            len(results),
            len(parsed),
            duration_ms,
        )
        return _success_result(
            {
                "title": f"批量画像完成 ({len(results)} 条)",
                "output": "\n".join(output_chunks),
                "results": results,
                "count": len(results),
                "duration_ms": duration_ms,
            }
        )

    except Exception as e:
        logger.error("batch_fetch_portraits unexpected error: error=%s", e, exc_info=True)
        return _error_result(f"未知错误: {e}", title="批量画像失败")


async def main() -> None:
    content_id = os.getenv("TEST_CONTENT_ID", "7641118685977614586")
    account_id = os.getenv("TEST_ACCOUNT_SEC_UID", "MS4wLjABAAAAcA9a--HmibvcoJ_0YCQYZ1qqbn2uCj5e4CVdc0c6y6s")

    print("=== 测试 get_content_fans_portrait ===")
    content_result = json.loads(await get_content_fans_portrait(content_id=content_id))
    if "error" in content_result:
        print(f"获取失败: {content_result['error']}")
    else:
        print(content_result["output"])
        print(
            f"\nhas_portrait={content_result.get('has_portrait')} "
            f"duration_ms={content_result.get('duration_ms')}"
        )

    if account_id:
        print("\n=== 测试 get_account_fans_portrait ===")
        account_result = json.loads(
            await get_account_fans_portrait(account_id=account_id)
        )
        if "error" in account_result:
            print(f"获取失败: {account_result['error']}")
        else:
            print(account_result["output"])
            print(
                f"\nhas_portrait={account_result.get('has_portrait')} "
                f"duration_ms={account_result.get('duration_ms')}"
            )
    else:
        print("\n跳过账号画像测试（设置环境变量 TEST_ACCOUNT_SEC_UID 可启用）")

    print("\n=== 测试 batch_fetch_portraits ===")
    candidates = [
        {
            "aweme_id": content_id,
            "author_sec_uid": account_id or None,
            "try_account_fallback": bool(account_id),
        }
    ]
    batch_result = json.loads(
        await batch_fetch_portraits(candidates_json=json.dumps(candidates, ensure_ascii=False))
    )
    if "error" in batch_result:
        print(f"批量获取失败: {batch_result['error']}")
    else:
        print(batch_result["output"])
        print(f"\ncount={batch_result.get('count')} duration_ms={batch_result.get('duration_ms')}")


if __name__ == "__main__":
    asyncio.run(main())