acquire-lead/db/repository.py

"""采集结果入库：videos / leads / comments。"""

from datetime import datetime

from sqlalchemy import select
from sqlalchemy.orm import Session

from config.settings import LEAD_STATUS_COLLECTED
from db.models import Comment, Lead, Video


def upsert_video(session: Session, meta: dict) -> Video:
    """按平台 video_id 插入或更新视频，返回 ORM 对象（含主键 id）。"""
    platform_video_id = (meta.get("video_id") or "").strip()
    if not platform_video_id:
        raise ValueError("video_id 为空，无法入库")

    now = datetime.now()
    video = session.scalar(select(Video).where(Video.video_id == platform_video_id))
    if video is None:
        video = Video(video_id=platform_video_id, create_at=now, update_at=now)
        session.add(video)
    else:
        video.update_at = now

    video.keyword = meta.get("keyword") or video.keyword
    video.title = meta.get("title") or video.title
    video.author_nickname = meta.get("author_nickname") or video.author_nickname
    video.author_douyin_id = meta.get("author_douyin_id") or video.author_douyin_id
    video.author_link = meta.get("author_link") or video.author_link
    video.video_url = meta.get("video_url") or video.video_url
    if meta.get("summary"):
        video.summary = str(meta["summary"])[:512]
    session.flush()
    return video


def upsert_lead(
    session: Session,
    *,
    douyin_id: str,
    nick_name: str | None,
    user_link: str | None,
) -> Lead:
    """按抖音号插入或更新线索；已存在时不覆盖更高阶 status。"""
    douyin_id = douyin_id.strip()
    if not douyin_id:
        raise ValueError("douyin_id 为空，无法入库")

    now = datetime.now()
    lead = session.scalar(select(Lead).where(Lead.douyin_id == douyin_id))
    if lead is None:
        lead = Lead(
            douyin_id=douyin_id,
            nick_name=nick_name,
            user_link=user_link,
            status=LEAD_STATUS_COLLECTED,
            create_at=now,
            update_at=now,
        )
        session.add(lead)
    else:
        if nick_name:
            lead.nick_name = nick_name
        if user_link:
            lead.user_link = user_link
        if lead.status is None:
            lead.status = LEAD_STATUS_COLLECTED
        lead.update_at = now
    session.flush()
    return lead


def insert_comment_if_new(
    session: Session,
    *,
    video_pk: int,
    lead_id: int | None,
    douyin_id: str | None,
    nick_name: str | None,
    content: str | None,
    comment_time: str | None,
    is_reply: bool = False,
) -> Comment | None:
    """
    写入评论。
    comments.video_id = videos.id（主键）。
    唯一性：同一视频主键 + 同一抖音号 + 同一正文。
    """
    if not video_pk:
        raise ValueError("video_pk 为空，评论无法入库")

    content_norm = (content or "").strip()
    douyin_norm = (douyin_id or "").strip() or None
    if not douyin_norm or not content_norm:
        return None

    exists = session.scalar(
        select(Comment.id).where(
            Comment.video_id == video_pk,
            Comment.douyin_id == douyin_norm,
            Comment.content == content_norm,
        )
    )
    if exists:
        return None

    comment = Comment(
        video_id=video_pk,
        lead_id=lead_id,
        douyin_id=douyin_norm,
        nick_name=nick_name,
        content=content_norm,
        comment_time=comment_time,
        type=0 if is_reply else 1,
        create_at=datetime.now(),
    )
    session.add(comment)
    session.flush()
    return comment


def save_video_only(session: Session, video_meta: dict) -> dict:
    """仅 upsert 视频；返回平台 video_id 与表主键 id。"""
    video = upsert_video(session, video_meta)
    return {"video_id": video.video_id, "id": video.id}


def save_comments_batch(session: Session, video_pk: int, comments: list[dict]) -> dict:
    """
    写入线索与评论。
    video_pk: videos.id
    用户唯一键：douyin_id
    评论唯一键：videos.id + douyin_id + content
    """
    if not video_pk:
        raise ValueError("video_pk 为空")

    lead_count = 0
    comment_count = 0
    skip_count = 0
    no_id_count = 0

    for item in comments:
        douyin_id = (item.get("douyin_id") or "").strip()
        if not douyin_id:
            no_id_count += 1
            skip_count += 1
            continue

        lead = upsert_lead(
            session,
            douyin_id=douyin_id,
            nick_name=item.get("nick_name"),
            user_link=item.get("user_link"),
        )
        lead_count += 1

        saved = insert_comment_if_new(
            session,
            video_pk=video_pk,
            lead_id=lead.id,
            douyin_id=douyin_id,
            nick_name=item.get("nick_name"),
            content=item.get("content"),
            comment_time=item.get("comment_time"),
            is_reply=bool(item.get("is_reply")),
        )
        if saved is None:
            skip_count += 1
        else:
            comment_count += 1

    return {
        "video_pk": video_pk,
        "leads": lead_count,
        "comments": comment_count,
        "skipped_comments": skip_count,
        "no_douyin_id": no_id_count,
    }


def save_collect_batch(session: Session, video_meta: dict, comments: list[dict]) -> dict:
    """兼容旧调用：一次事务写入视频 + 评论。"""
    video = upsert_video(session, video_meta)
    return save_comments_batch(session, video.id, comments)