Create-knowledge-find-decode/acquisition/creation_search.py

"""430 query 真实采集：搜索前 10、详情限速、媒体保存/OSS、创作知识判断。"""
from __future__ import annotations

import hashlib
import json
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Iterable, Optional

from acquisition import store
from acquisition.classify import classify_imgtext, classify_video
from acquisition.crawler import RateLimiter, fetch_post_detail
from acquisition.oss import upload_stream
from acquisition.search import search_keyword, search_weixin, search_xiaohongshu, fetch_weixin_detail
from core.config import Settings
from core.models import Post
from creation_knowledge.integrations.video_extract import _default_download

ROOT = Path(__file__).resolve().parent.parent
DEFAULT_QUERY_FILE = ROOT / "data" / "queries" / "creation_demo.json"
PLATFORMS = ("xiaohongshu", "weixin", "douyin")

Downloader = Callable[[str, str], bytes]


class PlatformRateLimiter:
    """把同一平台的 keyword/detail 调用合并到同一个 10-12s bucket。"""

    def __init__(self, platform: str, *, min_seconds: float = 10.0, max_seconds: float = 12.0,
                 delegate: Optional[RateLimiter] = None) -> None:
        self.platform = platform
        self.delegate = delegate or RateLimiter(
            min_interval_seconds=min_seconds,
            max_interval_seconds=max_seconds,
        )

    def wait(self, bucket: str) -> None:
        self.delegate.wait(self.platform)


@dataclass
class Candidate:
    rank: int
    source_id: str = ""
    url: str = ""
    title: str = ""
    author: str = ""
    cover_url: str = ""
    raw: dict | None = None


def load_creation_queries(path: Path | str = DEFAULT_QUERY_FILE) -> list[str]:
    """从 creation_demo.json 读取所有 keep=true 的唯一 query，保持首次出现顺序。"""
    data = json.loads(Path(path).read_text("utf-8"))
    seen: set[str] = set()
    out: list[str] = []
    for fam in data.get("families") or []:
        for item in fam.get("items") or []:
            q = (item.get("query") or "").strip()
            if q and item.get("keep", True) and q not in seen:
                seen.add(q)
                out.append(q)
    return out


def prompt_version(*names: str) -> str:
    h = hashlib.sha256()
    for name in names:
        p = ROOT / "prompts" / f"{name}.txt"
        h.update(name.encode("utf-8"))
        if p.exists():
            h.update(p.read_bytes())
    return h.hexdigest()[:16]


def _hash(value: str, n: int = 16) -> str:
    return hashlib.md5(value.encode("utf-8")).hexdigest()[:n]


def _ext_from_bytes(data: bytes) -> str:
    if data[:4] == b"\x89PNG":
        return ".png"
    if data[:3] == b"GIF":
        return ".gif"
    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        return ".webp"
    return ".jpg"


def _dedup(values: Iterable[str]) -> list[str]:
    seen: set[str] = set()
    out: list[str] = []
    for v in values:
        if v and v not in seen:
            seen.add(v)
            out.append(v)
    return out


def _data_root(settings: Settings) -> Path:
    p = Path(settings.data_dir or "data")
    return p if p.is_absolute() else ROOT / p


def _save_images(urls: list[str], *, platform: str, base_dir: Path, public_base: str,
                 downloader: Downloader = _default_download, max_images: int = 12) -> list[str]:
    out: list[str] = []
    for i, url in enumerate(_dedup(urls)[:max_images], start=1):
        if not url.startswith("http"):
            continue
        try:
            data = downloader(url, platform)
        except Exception:
            continue
        if not data:
            continue
        base_dir.mkdir(parents=True, exist_ok=True)
        ext = _ext_from_bytes(data)
        name = f"image_{i}{ext}"
        (base_dir / name).write_bytes(data)
        out.append(f"{public_base}/{name}")
    return out


def _classify_and_store(conn, item_id: int, platform: str, *, title: str, body_text: str,
                        images: list[str], video_url: str, settings: Settings,
                        ts: Optional[int] = None) -> None:
    if platform == "douyin":
        version = prompt_version("classify_video")
        is_creation, reason, knowledge, _points = classify_video(
            {"platform": platform, "title": title, "body_text": body_text, "video": video_url},
            settings,
        )
    else:
        version = prompt_version("classify_imgtext")
        is_creation, reason, knowledge, _points = classify_imgtext(
            {"platform": platform, "title": title, "body_text": body_text, "images": images},
            settings,
        )
    store.upsert_creation_classification(
        conn,
        item_id,
        is_creation,
        reason=reason,
        knowledge=knowledge,
        prompt_version=version,
        error=reason if is_creation is None else "",
        ts=ts,
    )


def search_candidates(platform: str, query: str, *, settings: Settings, limit: int,
                      rate_limiter: PlatformRateLimiter) -> list[Candidate]:
    if platform == "xiaohongshu":
        rows = search_xiaohongshu(
            query, content_type="图文", limit=limit, settings=settings, rate_limiter=rate_limiter
        )
        return [
            Candidate(
                rank=i, source_id=r.get("id") or "", url=r.get("url") or "",
                title=r.get("title") or "", author=r.get("nick_name") or "",
                cover_url=r.get("cover_url") or "", raw=r,
            )
            for i, r in enumerate(rows, start=1)
        ]
    if platform == "weixin":
        rows = search_weixin(query, limit=limit, settings=settings, rate_limiter=rate_limiter)
        return [
            Candidate(
                rank=i, url=r.get("url") or "", title=r.get("title") or "",
                author=r.get("nick_name") or "", cover_url=r.get("cover_url") or "", raw=r,
            )
            for i, r in enumerate(rows, start=1)
        ]
    if platform == "douyin":
        ids = search_keyword(
            query, platform="douyin", content_type="视频", limit=limit,
            settings=settings, rate_limiter=rate_limiter,
        )
        return [Candidate(rank=i, source_id=cid, raw={"id": cid}) for i, cid in enumerate(ids, start=1)]
    raise ValueError(f"unsupported platform: {platform}")


def _process_xhs(candidate: Candidate, *, query_hash: str, settings: Settings,
                 rate_limiter: PlatformRateLimiter, downloader: Downloader) -> dict:
    post = fetch_post_detail(candidate.source_id or candidate.url, settings=settings, rate_limiter=rate_limiter)
    base = _data_root(settings) / "media" / "xiaohongshu" / query_hash / (post.content_id or candidate.source_id)
    pub = f"/data/media/xiaohongshu/{query_hash}/{post.content_id or candidate.source_id}"
    images = _save_images(post.image_urls or [candidate.cover_url], platform="xiaohongshu",
                          base_dir=base, public_base=pub, downloader=downloader)
    if not images:
        raise RuntimeError("图片下载失败")
    return {
        "source_id": post.content_id or candidate.source_id,
        "url": post.url or candidate.url,
        "title": post.title or candidate.title,
        "author": post.author_name or candidate.author,
        "body_text": post.body_text or "",
        "cover_url": images[0],
        "image_urls": images,
        "video_url": "",
        "raw": post.raw if isinstance(post.raw, dict) else {},
    }


def _process_weixin(candidate: Candidate, *, query_hash: str, settings: Settings,
                    rate_limiter: PlatformRateLimiter, downloader: Downloader) -> dict:
    body_text, detail_images = fetch_weixin_detail(
        candidate.url, settings=settings, rate_limiter=rate_limiter
    )
    h = _hash(candidate.url)
    base = _data_root(settings) / "media" / "weixin" / query_hash / h
    pub = f"/data/media/weixin/{query_hash}/{h}"
    images = _save_images(detail_images or [candidate.cover_url], platform="weixin",
                          base_dir=base, public_base=pub, downloader=downloader)
    return {
        "source_id": h,
        "url": candidate.url,
        "title": candidate.title,
        "author": candidate.author,
        "body_text": body_text,
        "cover_url": images[0] if images else candidate.cover_url,
        "image_urls": images,
        "video_url": "",
        "raw": candidate.raw or {},
    }


def _process_douyin(candidate: Candidate, *, settings: Settings,
                    rate_limiter: PlatformRateLimiter) -> dict:
    post: Post = fetch_post_detail(candidate.source_id, settings=settings, rate_limiter=rate_limiter)
    if not post.video_urls:
        raise RuntimeError("详情无视频")
    cdn_url = upload_stream(post.video_urls[0], src_type="video", settings=settings)
    return {
        "source_id": post.content_id or candidate.source_id,
        "url": post.url,
        "title": post.title,
        "author": post.author_name or "",
        "body_text": post.body_text or "",
        "cover_url": (post.image_urls or [""])[0],
        "image_urls": post.image_urls or [],
        "video_url": cdn_url,
        "raw": post.raw if isinstance(post.raw, dict) else {},
    }


def process_candidate(platform: str, candidate: Candidate, *, query_hash: str,
                      settings: Settings, rate_limiter: PlatformRateLimiter,
                      downloader: Downloader = _default_download) -> dict:
    if platform == "xiaohongshu":
        return _process_xhs(candidate, query_hash=query_hash, settings=settings,
                            rate_limiter=rate_limiter, downloader=downloader)
    if platform == "weixin":
        return _process_weixin(candidate, query_hash=query_hash, settings=settings,
                               rate_limiter=rate_limiter, downloader=downloader)
    if platform == "douyin":
        return _process_douyin(candidate, settings=settings, rate_limiter=rate_limiter)
    raise ValueError(f"unsupported platform: {platform}")


def run_platform_query(conn, *, run_id: str, query: str, platform: str, settings: Settings,
                       search_limit: int = 10, display_limit: int = 5,
                       rate_limiter: Optional[PlatformRateLimiter] = None,
                       downloader: Downloader = _default_download,
                       sleep_fn: Callable[[float], None] = time.sleep,
                       classify: bool = True, ts: Optional[int] = None) -> dict:
    """执行一个 query×platform job。任何 item 失败都会落库并继续。"""
    gate = rate_limiter or PlatformRateLimiter(platform)
    now = int(time.time()) if ts is None else ts
    store.update_creation_job(conn, run_id, query, platform, status="running", ts=now)

    candidates: list[Candidate] = []
    last_error = ""
    attempts = 0
    for attempts in range(1, 4):
        try:
            candidates = search_candidates(
                platform, query, settings=settings, limit=search_limit, rate_limiter=gate
            )
            last_error = ""
            break
        except Exception as exc:
            last_error = f"搜索失败: {str(exc)[:120]}"
            store.update_creation_job(
                conn, run_id, query, platform, status="running", attempts=attempts,
                error=last_error, ts=int(time.time()),
            )
            if attempts < 3:
                sleep_fn(30 if attempts == 1 else 90)

    if last_error:
        store.update_creation_job(
            conn, run_id, query, platform, status="failed", attempts=attempts,
            searched_count=0, display_count=0, error=last_error, ts=int(time.time()),
        )
        return {"platform": platform, "status": "failed", "display_count": 0, "error": last_error}

    qh = _hash(query)
    display_count = 0
    errors: list[str] = []
    for cand in candidates[:search_limit]:
        if display_count >= display_limit:
            break
        try:
            processed = process_candidate(
                platform, cand, query_hash=qh, settings=settings, rate_limiter=gate,
                downloader=downloader,
            )
            item_id = store.upsert_creation_item(
                conn, run_id=run_id, query=query, platform=platform, rank=cand.rank,
                source_id=processed.get("source_id", ""), url=processed.get("url", ""),
                title=processed.get("title", ""), author=processed.get("author", ""),
                body_text=processed.get("body_text", ""), cover_url=processed.get("cover_url", ""),
                image_urls=processed.get("image_urls", []), video_url=processed.get("video_url", ""),
                media={"source": cand.raw or {}, "candidate_rank": cand.rank},
                raw=processed.get("raw", {}), is_displayable=True, ts=int(time.time()),
            )
            if classify:
                _classify_and_store(
                    conn, item_id, platform, title=processed.get("title", ""),
                    body_text=processed.get("body_text", ""),
                    images=processed.get("image_urls", []),
                    video_url=processed.get("video_url", ""), settings=settings,
                    ts=int(time.time()),
                )
            display_count += 1
        except Exception as exc:
            msg = f"rank {cand.rank}: {str(exc)[:120]}"
            errors.append(msg)
            store.upsert_creation_item(
                conn, run_id=run_id, query=query, platform=platform, rank=cand.rank,
                source_id=cand.source_id, url=cand.url, title=cand.title,
                author=cand.author, cover_url=cand.cover_url, raw=cand.raw or {},
                is_displayable=False, error=msg, ts=int(time.time()),
            )

    if display_count >= display_limit:
        status = "done"
        error = ""
    elif display_count > 0:
        status = "partial"
        error = "; ".join(errors[-3:])
    else:
        status = "failed"
        error = "; ".join(errors[-3:]) or "未获得可展示内容"
    store.update_creation_job(
        conn, run_id, query, platform, status=status, attempts=attempts,
        searched_count=len(candidates), display_count=display_count,
        error=error, ts=int(time.time()),
    )
    return {"platform": platform, "status": status, "display_count": display_count, "error": error}