Agent/examples/auto_put_ad_mini/tools/external_recalled_material.py

"""External recalled material processing for module B.

This path is enabled only for accounts configured with
``material_source=external_recall``. It keeps historical recall and the
existing text-to-image flow unchanged.
"""

from __future__ import annotations

import hashlib
import json
import logging
import os
from dataclasses import dataclass
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Iterable

from tools.ai_generated_material import (
    OPENROUTER_IMAGE_MODEL,
    build_ai_image_object_key,
    generate_image_bytes,
    upload_image_to_oss,
)
from tools.ai_material_review import MaterialReviewResult, review_generated_material
from tools.material_recall import Material, recall_materials_for_video
from tools.video_recall import LandingVideo

logger = logging.getLogger(__name__)

EXTERNAL_RECALL_SOURCE_LABEL = os.getenv(
    "EXTERNAL_RECALL_SOURCE_LABEL", "外部合作"
).strip()
EXTERNAL_RECALL_CANDIDATE_LIMIT = int(
    os.getenv("EXTERNAL_RECALL_CANDIDATE_LIMIT", "300")
)
EXTERNAL_RECALL_EDIT_LIMIT_PER_LANDING = int(
    os.getenv("EXTERNAL_RECALL_EDIT_LIMIT_PER_LANDING", "3")
)
EXTERNAL_RECALL_UV_WINDOW_DAYS = int(
    os.getenv("EXTERNAL_RECALL_UV_WINDOW_DAYS", "90")
)
EXTERNAL_IMAGE_MODEL = os.getenv(
    "EXTERNAL_IMAGE_MODEL", OPENROUTER_IMAGE_MODEL
).strip()
EXTERNAL_CLEANUP_PROMPT_PATH = (
    Path(__file__).resolve().parents[1] / "prompts" / "external_material_cleanup.md"
)


CREATE_EXTERNAL_MATERIAL_TABLE_SQL = """
CREATE TABLE IF NOT EXISTS external_recalled_material (
    id BIGINT AUTO_INCREMENT PRIMARY KEY COMMENT '自增主键',
    account_id BIGINT NOT NULL COMMENT '腾讯广告账户ID',
    adgroup_id BIGINT NOT NULL COMMENT '广告ID',
    crowd_package VARCHAR(200) NOT NULL COMMENT '投放人群包名称',
    landing_video_id BIGINT NOT NULL COMMENT '承接视频ID',
    source_material_id VARCHAR(255) NOT NULL COMMENT '外部召回materialId/card_cover_id',
    source_image_url VARCHAR(1000) NOT NULL COMMENT '外部素材原图URL',
    source_title VARCHAR(1000) DEFAULT NULL COMMENT '外部素材标题',
    similarity_score DOUBLE NOT NULL DEFAULT 0 COMMENT '向量召回相似度',
    visit_uv_30d BIGINT NOT NULL DEFAULT 0 COMMENT '配置窗口内访问UV之和(兼容历史字段名)',
    uv_window_start VARCHAR(8) DEFAULT NULL COMMENT 'UV窗口开始分区',
    uv_window_end VARCHAR(8) DEFAULT NULL COMMENT 'UV窗口结束分区',
    recall_json MEDIUMTEXT DEFAULT NULL COMMENT '召回命中与原始素材JSON',
    edit_prompt MEDIUMTEXT NOT NULL COMMENT '去播放按钮图生图prompt',
    model VARCHAR(200) NOT NULL COMMENT '图生图模型',
    oss_object_key VARCHAR(500) DEFAULT NULL COMMENT '派生图OSS key',
    oss_url VARCHAR(1000) DEFAULT NULL COMMENT '派生图公网URL',
    status VARCHAR(50) NOT NULL DEFAULT 'generated' COMMENT 'generated/prepared/approved/rejected/hold/skip/posted_ok/post_failed/error',
    approval_status VARCHAR(50) DEFAULT NULL COMMENT '人工审批状态',
    tencent_image_id VARCHAR(100) DEFAULT NULL COMMENT '腾讯图片ID',
    dynamic_creative_id BIGINT DEFAULT NULL COMMENT '腾讯动态创意ID',
    ai_review_status VARCHAR(50) DEFAULT NULL COMMENT 'AI审核状态',
    ai_review_score INT DEFAULT NULL COMMENT 'AI审核评分',
    ai_review_reason TEXT DEFAULT NULL COMMENT 'AI审核原因',
    ai_review_json MEDIUMTEXT DEFAULT NULL COMMENT 'AI审核JSON',
    error TEXT DEFAULT NULL COMMENT '处理错误',
    raw_response MEDIUMTEXT DEFAULT NULL COMMENT '图生图模型响应摘要',
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
    UNIQUE KEY uk_account_ad_landing_source
        (account_id, adgroup_id, landing_video_id, source_material_id),
    KEY idx_account_ad_status (account_id, adgroup_id, status),
    KEY idx_source_created (source_material_id, created_at),
    KEY idx_status_created (status, created_at)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='外部召回素材AI清理派生图'
"""


@dataclass(frozen=True)
class ExternalMaterialAsset:
    id: int
    source_material_id: str
    source_image_url: str
    source_title: str
    similarity_score: float
    visit_uv_30d: int
    uv_window_start: str
    uv_window_end: str
    oss_url: str
    recall_json: dict[str, Any]

    def to_material(self) -> Material:
        recall = self.recall_json
        return Material(
            material_id=f"external:{self.source_material_id}",
            score=self.similarity_score,
            title=self.source_title,
            cover=self.oss_url,
            cost=None,
            ctr=None,
            cvr=None,
            roi=None,
            impressions=None,
            quality_score=None,
            visit_uv_30d=self.visit_uv_30d,
            recall_strategy=str(recall.get("recall_strategy") or ""),
            recall_query_text=str(recall.get("recall_query_text") or ""),
            recall_config_code=str(recall.get("recall_config_code") or ""),
            recall_element_dimension=str(
                recall.get("recall_element_dimension") or ""
            ),
            recall_point_type=str(recall.get("recall_point_type") or ""),
            recall_standard_element=str(
                recall.get("recall_standard_element") or ""
            ),
            recall_hit_queries=list(recall.get("recall_hit_queries") or []),
            raw={
                "external_recalled_material_id": self.id,
                "source_material_id": self.source_material_id,
                "source_image_url": self.source_image_url,
                "visit_uv_30d": self.visit_uv_30d,
                "uv_window_start": self.uv_window_start,
                "uv_window_end": self.uv_window_end,
                "oss_url": self.oss_url,
            },
        )


def ensure_external_material_table() -> None:
    from db.connection import get_connection

    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(CREATE_EXTERNAL_MATERIAL_TABLE_SQL)
        conn.commit()
    finally:
        conn.close()


def _source_image_url(material: Material) -> str:
    if material.cover:
        return material.cover
    images = (material.raw or {}).get("imageList") or []
    return str(images[0]) if images else ""


def _quote_sql(value: str) -> str:
    return "'" + value.replace("\\", "\\\\").replace("'", "\\'") + "'"


def _get_odps_client():
    from tools.odps_module import get_odps_client

    client = get_odps_client(project="loghubods")
    if client is None:
        raise RuntimeError("无法初始化ODPS客户端")
    return client


def _latest_partition_dt(client) -> str:
    table = client._odps.get_table("cooperate_top_cards_daily_v1")
    partition = table.get_max_partition()
    if partition is None:
        raise RuntimeError("cooperate_top_cards_daily_v1没有可用分区")
    name = getattr(partition, "name", "") or str(partition)
    for part in name.split(","):
        if part.strip().startswith("dt="):
            return part.split("=", 1)[1].strip(" '\"")
    raise RuntimeError(f"无法解析cooperate_top_cards_daily_v1最新分区:{name}")


def _query_visit_uv_30d(
    material_ids: Iterable[str],
) -> tuple[dict[str, int], str, str]:
    ids = sorted({str(value).strip() for value in material_ids if str(value).strip()})
    if not ids:
        return {}, "", ""
    client = _get_odps_client()
    end_dt = _latest_partition_dt(client)
    end_date = datetime.strptime(end_dt, "%Y%m%d").date()
    start_dt = (end_date - timedelta(days=EXTERNAL_RECALL_UV_WINDOW_DAYS - 1)).strftime(
        "%Y%m%d"
    )
    totals: dict[str, int] = {}
    for offset in range(0, len(ids), 500):
        batch = ids[offset:offset + 500]
        quoted = ",".join(_quote_sql(value) for value in batch)
        sql = f"""
SELECT card_cover_id, SUM(`访问uv`) AS visit_uv_30d
FROM loghubods.cooperate_top_cards_daily_v1
WHERE dt BETWEEN '{start_dt}' AND '{end_dt}'
  AND card_cover_id IN ({quoted})
GROUP BY card_cover_id
"""
        instance = client._odps.execute_sql(sql)
        instance.wait_for_success()
        with instance.open_reader(tunnel=False) as reader:
            for row in reader:
                totals[str(row[0])] = int(row[1] or 0)
    return totals, start_dt, end_dt


def rank_external_materials_by_visit_uv(
    materials: list[Material],
) -> list[Material]:
    """Attach configured-window visit UV and sort by UV, then similarity."""
    if not materials:
        return []
    totals, start_dt, end_dt = _query_visit_uv_30d(
        material.material_id for material in materials
    )
    for material in materials:
        material.visit_uv_30d = totals.get(material.material_id, 0)
        material.raw["visit_uv_30d"] = material.visit_uv_30d
        material.raw["uv_window_start"] = start_dt
        material.raw["uv_window_end"] = end_dt
    ranked = sorted(
        materials,
        key=lambda material: (material.visit_uv_30d or 0, material.score or 0),
        reverse=True,
    )
    logger.info(
        "[external_material] UV排序 candidates=%d days=%d window=%s-%s top=%s",
        len(ranked), EXTERNAL_RECALL_UV_WINDOW_DAYS, start_dt, end_dt,
        [
            (material.material_id, material.visit_uv_30d, round(material.score, 4))
            for material in ranked[:5]
        ],
    )
    return ranked


def recall_external_materials_for_video(
    landing: LandingVideo,
    *,
    element_features: Iterable,
) -> list[Material]:
    materials = recall_materials_for_video(
        landing,
        final_top_n=None,
        source_labels=[EXTERNAL_RECALL_SOURCE_LABEL],
        element_features=element_features,
        apply_cover_blacklist=False,
    )
    ranked = rank_external_materials_by_visit_uv(materials)
    return ranked[:max(1, EXTERNAL_RECALL_CANDIDATE_LIMIT)]


def _load_cleanup_prompt() -> str:
    text = EXTERNAL_CLEANUP_PROMPT_PATH.read_text(encoding="utf-8").strip()
    if not text:
        raise RuntimeError(f"外部素材清理prompt为空:{EXTERNAL_CLEANUP_PROMPT_PATH}")
    return text


def _recall_json(material: Material) -> dict[str, Any]:
    return {
        "recall_strategy": material.recall_strategy,
        "recall_query_text": material.recall_query_text,
        "recall_config_code": material.recall_config_code,
        "recall_element_dimension": material.recall_element_dimension,
        "recall_point_type": material.recall_point_type,
        "recall_standard_element": material.recall_standard_element,
        "recall_hit_queries": material.recall_hit_queries,
        "source_raw": material.raw,
    }


def _asset_from_row(row: dict[str, Any]) -> ExternalMaterialAsset:
    try:
        recall = json.loads(row.get("recall_json") or "{}")
    except Exception:
        recall = {}
    return ExternalMaterialAsset(
        id=int(row["id"]),
        source_material_id=str(row["source_material_id"]),
        source_image_url=str(row.get("source_image_url") or ""),
        source_title=str(row.get("source_title") or ""),
        similarity_score=float(row.get("similarity_score") or 0),
        visit_uv_30d=int(row.get("visit_uv_30d") or 0),
        uv_window_start=str(row.get("uv_window_start") or ""),
        uv_window_end=str(row.get("uv_window_end") or ""),
        oss_url=str(row.get("oss_url") or ""),
        recall_json=recall,
    )


def _load_existing_row(
    account_id: int,
    adgroup_id: int,
    landing_video_id: int,
    source_material_id: str,
) -> dict[str, Any] | None:
    ensure_external_material_table()
    from db.connection import get_connection

    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
SELECT *
FROM external_recalled_material
WHERE account_id=%s AND adgroup_id=%s AND landing_video_id=%s
  AND source_material_id=%s
LIMIT 1
""",
                (account_id, adgroup_id, landing_video_id, source_material_id),
            )
            return cur.fetchone()
    finally:
        conn.close()


def _insert_generated_asset(
    *,
    account_id: int,
    adgroup_id: int,
    crowd_package: str,
    landing: LandingVideo,
    material: Material,
    source_image_url: str,
    prompt: str,
    object_key: str,
    oss_url: str,
    raw_response: dict[str, Any],
) -> int:
    ensure_external_material_table()
    from db.connection import get_connection

    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
INSERT INTO external_recalled_material
    (account_id, adgroup_id, crowd_package, landing_video_id,
     source_material_id, source_image_url, source_title, similarity_score,
     visit_uv_30d, uv_window_start, uv_window_end, recall_json,
     edit_prompt, model, oss_object_key, oss_url, raw_response)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
""",
                (
                    account_id,
                    adgroup_id,
                    crowd_package,
                    landing.video_id,
                    material.material_id,
                    source_image_url,
                    material.title,
                    material.score,
                    int(material.visit_uv_30d or 0),
                    str(material.raw.get("uv_window_start") or ""),
                    str(material.raw.get("uv_window_end") or ""),
                    json.dumps(_recall_json(material), ensure_ascii=False, default=str),
                    prompt,
                    EXTERNAL_IMAGE_MODEL,
                    object_key,
                    oss_url,
                    json.dumps(raw_response, ensure_ascii=False, default=str)[:16000000],
                ),
            )
            asset_id = int(cur.lastrowid)
        conn.commit()
        return asset_id
    finally:
        conn.close()


def _update_review(asset_id: int, review) -> None:
    from db.connection import get_connection

    status = "generated" if review.status == "pass" else review.status
    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
UPDATE external_recalled_material
SET status=%s, ai_review_status=%s, ai_review_score=%s,
    ai_review_reason=%s, ai_review_json=%s, error=%s,
    updated_at=CURRENT_TIMESTAMP
WHERE id=%s
""",
                (
                    status,
                    review.status,
                    review.score,
                    review.reason,
                    json.dumps(review.raw, ensure_ascii=False, default=str),
                    review.reason if review.status == "error" else None,
                    asset_id,
                ),
            )
        conn.commit()
    finally:
        conn.close()


def _review_asset(
    *,
    asset_id: int,
    oss_url: str,
    prompt: str,
    landing: LandingVideo,
    material: Material,
):
    try:
        review = review_generated_material(
            image_url=oss_url,
            prompt_type="external_cleanup",
            prompt_text=prompt,
            feature_hits=[
                {
                    "landing_video_id": landing.video_id,
                    "landing_title": landing.title,
                    "source_material_id": material.material_id,
                    "source_title": material.title,
                    "visit_uv_30d": int(material.visit_uv_30d or 0),
                    "recall_hit_queries": material.recall_hit_queries,
                }
            ],
        )
    except Exception as exc:
        logger.exception(
            "[external_material] AI审核异常 asset=%d source=%s:%s",
            asset_id, material.material_id, exc,
        )
        review = MaterialReviewResult(
            status="error",
            score=0,
            reason=f"AI审核异常:{exc}",
            risk_tags=["review_error"],
            ocr_text="",
            raw={"error": str(exc)},
        )
    _update_review(asset_id, review)
    return review


def get_or_create_external_asset(
    *,
    account_id: int,
    adgroup_id: int,
    crowd_package: str,
    landing: LandingVideo,
    material: Material,
) -> Material | None:
    existing = _load_existing_row(
        account_id,
        adgroup_id,
        landing.video_id,
        material.material_id,
    )
    if existing:
        if (
            existing.get("oss_url")
            and existing.get("status") in {"generated", "error"}
            and existing.get("ai_review_status") in {None, "", "error"}
        ):
            review = _review_asset(
                asset_id=int(existing["id"]),
                oss_url=str(existing["oss_url"]),
                prompt=str(existing.get("edit_prompt") or _load_cleanup_prompt()),
                landing=landing,
                material=material,
            )
            existing = _load_existing_row(
                account_id,
                adgroup_id,
                landing.video_id,
                material.material_id,
            )
            if review.passed and existing:
                return _asset_from_row(existing).to_material()
        if (
            existing.get("ai_review_status") == "pass"
            and existing.get("status") in {"generated", "prepared"}
            and existing.get("oss_url")
        ):
            logger.info(
                "[external_material] 复用派生图 account=%d adgroup=%d landing=%d source=%s id=%s",
                account_id, adgroup_id, landing.video_id, material.material_id,
                existing["id"],
            )
            return _asset_from_row(existing).to_material()
        logger.info(
            "[external_material] 已有不可用派生记录 source=%s status=%s review=%s",
            material.material_id, existing.get("status"),
            existing.get("ai_review_status"),
        )
        return None

    source_image_url = _source_image_url(material)
    if not source_image_url:
        logger.warning(
            "[external_material] source=%s 没有可编辑图片URL", material.material_id
        )
        return None

    prompt = _load_cleanup_prompt()
    image_bytes, content_type, raw_response = generate_image_bytes(
        prompt,
        model=EXTERNAL_IMAGE_MODEL,
        input_image_url=source_image_url,
    )
    source_hash = hashlib.sha256(material.material_id.encode("utf-8")).hexdigest()[:12]
    object_key = build_ai_image_object_key(
        account_id=account_id,
        landing_video_id=landing.video_id,
        prompt_type=f"external_cleanup_{source_hash}",
        extension="jpg",
    )
    oss_url = upload_image_to_oss(image_bytes, content_type, object_key)
    asset_id = _insert_generated_asset(
        account_id=account_id,
        adgroup_id=adgroup_id,
        crowd_package=crowd_package,
        landing=landing,
        material=material,
        source_image_url=source_image_url,
        prompt=prompt,
        object_key=object_key,
        oss_url=oss_url,
        raw_response=raw_response,
    )
    review = _review_asset(
        asset_id=asset_id,
        oss_url=oss_url,
        prompt=prompt,
        landing=landing,
        material=material,
    )
    logger.info(
        "[external_material] 派生图完成 account=%d adgroup=%d landing=%d source=%s "
        "asset=%d uv_window=%d review=%s score=%d url=%s",
        account_id, adgroup_id, landing.video_id, material.material_id,
        asset_id, int(material.visit_uv_30d or 0), review.status, review.score, oss_url,
    )
    if not review.passed:
        return None
    row = _load_existing_row(
        account_id,
        adgroup_id,
        landing.video_id,
        material.material_id,
    )
    return _asset_from_row(row).to_material() if row else None


def prepare_external_material_for_landing(
    *,
    account_id: int,
    adgroup_id: int,
    crowd_package: str,
    landing: LandingVideo,
    element_features: Iterable,
    excluded_material_ids: set[str],
) -> Material | None:
    ranked = recall_external_materials_for_video(
        landing,
        element_features=element_features,
    )
    attempted = 0
    for material in ranked:
        external_id = f"external:{material.material_id}"
        if external_id in excluded_material_ids:
            continue
        attempted += 1
        try:
            processed = get_or_create_external_asset(
                account_id=account_id,
                adgroup_id=adgroup_id,
                crowd_package=crowd_package,
                landing=landing,
                material=material,
            )
        except Exception as exc:
            logger.exception(
                "[external_material] 图生图失败 landing=%d source=%s:%s",
                landing.video_id, material.material_id, exc,
            )
            processed = None
        if processed is not None:
            return processed
        if attempted >= max(1, EXTERNAL_RECALL_EDIT_LIMIT_PER_LANDING):
            break
    logger.info(
        "[external_material] landing=%d 无可用派生素材 recalled=%d attempted=%d",
        landing.video_id, len(ranked), attempted,
    )
    return None


def update_external_material_status(
    record: dict,
    status: str,
    *,
    dynamic_creative_id: int | str | None = None,
    tencent_image_id: str = "",
    error: str = "",
) -> None:
    asset_id = record.get("_external_recalled_material_id")
    if not asset_id:
        return
    ensure_external_material_table()
    from db.connection import get_connection

    normalized_status = {
        "approve": "approved",
        "reject": "rejected",
    }.get(status, status)
    approval_status = {
        "approve": "approved",
        "reject": "rejected",
        "hold": "hold",
        "skip": "skip",
    }.get(status, normalized_status)
    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
UPDATE external_recalled_material
SET status=%s, approval_status=%s,
    dynamic_creative_id=COALESCE(%s, dynamic_creative_id),
    tencent_image_id=COALESCE(NULLIF(%s, ''), tencent_image_id),
    error=COALESCE(NULLIF(%s, ''), error), updated_at=CURRENT_TIMESTAMP
WHERE id=%s
""",
                (
                    normalized_status,
                    approval_status,
                    int(dynamic_creative_id) if dynamic_creative_id else None,
                    tencent_image_id,
                    error[:2000] if error else "",
                    int(asset_id),
                ),
            )
        conn.commit()
    finally:
        conn.close()