life-echo/api/app/tasks/chapter_cover_tasks.py

"""
Chapter 封面生成 Celery 任务。

从 chapter_cover_intents 原子 claim intent，或创建新 intent 后生成封面，
写入 assets，绑定到 chapters.cover_asset_id。封面不回写进正文 markdown。
"""

import hashlib
import time
import uuid
from datetime import datetime, timedelta, timezone

from celery import shared_task
from PIL import Image
from sqlalchemy import and_, func, or_, select, update
from sqlalchemy.orm import joinedload

from app.agents.image_prompt import get_image_prompt_orchestrator
from app.core.db import get_sync_db, transactional_sync
from app.core.dependencies import get_image_generator
from app.core.logging import get_logger
from app.core.redis_lock import acquire_redis_lock, release_redis_lock
from app.features.asset.models import Asset
from app.features.memoir.chapter_cover import (
    aggregate_cover_prompt_from_stories,
)
from app.features.memoir.cover_eligibility import (
    chapter_eligible_for_cover_by_inline_body_image_count,
    chapter_has_story_links,
)
from app.features.memoir.memoir_images.storage import TencentCosStorageService
from app.features.memoir.models import Chapter, ChapterCoverIntent, ChapterStoryLink
from app.ports.image_gen import TaskStatus

logger = get_logger(__name__)

CHAPTER_COVER_LOCK_TTL_SECONDS = 1800
CHAPTER_COVER_CLAIM_TTL_SECONDS = 1800


class _ClaimSkipped(Exception):
    """Concurrent worker won the intent claim; abort transactional block."""


def _build_cover_cos_key(user_id: str, chapter_id: str, prompt: str) -> str:
    short_hash = hashlib.sha1(prompt.encode("utf-8")).hexdigest()[:10]
    return f"chapters/{user_id}/{chapter_id}/cover-{short_hash}.png"


def _normalize_image_bytes(image_bytes: bytes) -> bytes:
    from io import BytesIO

    with Image.open(BytesIO(image_bytes)) as image:
        output = BytesIO()
        if image.mode in {"RGBA", "LA"}:
            normalized = image
        elif image.mode == "P":
            normalized = image.convert("RGBA")
        else:
            normalized = image.convert("RGB")
        normalized.save(output, format="PNG")
        return output.getvalue()


def _chapter_cover_claimable_clause(now: datetime):
    cutoff = now - timedelta(seconds=CHAPTER_COVER_CLAIM_TTL_SECONDS)
    return or_(
        ChapterCoverIntent.status.in_(["pending", "failed"]),
        and_(
            ChapterCoverIntent.status == "processing",
            or_(
                ChapterCoverIntent.claimed_at.is_(None),
                ChapterCoverIntent.claimed_at < cutoff,
            ),
        ),
    )


def _build_chapter_cover_brief(chapter: Chapter) -> str:
    stories = []
    for link in sorted(
        chapter.story_links, key=lambda link_row: link_row.order_index or 0
    ):
        story = getattr(link, "story", None)
        if story:
            stories.append(story)
    return aggregate_cover_prompt_from_stories(
        stories,
        chapter_title=chapter.title or "",
        chapter_category=chapter.category or "",
    )


def _claim_chapter_cover_intent_sync(db, chapter: Chapter, claim_token: str):
    now = datetime.now(timezone.utc)
    claimable = _chapter_cover_claimable_clause(now)
    candidate_id = db.execute(
        select(ChapterCoverIntent.id)
        .where(ChapterCoverIntent.chapter_id == chapter.id)
        .where(claimable)
        .order_by(
            ChapterCoverIntent.updated_at.desc(), ChapterCoverIntent.created_at.desc()
        )
        .limit(1)
    ).scalar_one_or_none()
    if candidate_id:
        try:
            with transactional_sync(db):
                claimed = db.execute(
                    update(ChapterCoverIntent)
                    .where(ChapterCoverIntent.id == candidate_id)
                    .where(_chapter_cover_claimable_clause(now))
                    .values(
                        status="processing",
                        claim_token=claim_token,
                        claimed_at=now,
                        updated_at=now,
                        error=None,
                        attempt_count=func.coalesce(
                            ChapterCoverIntent.attempt_count, 0
                        )
                        + 1,
                    )
                )
                if (claimed.rowcount or 0) != 1:
                    raise _ClaimSkipped()
                intent = db.get(ChapterCoverIntent, candidate_id)
        except _ClaimSkipped:
            return None
        return intent

    cutoff = now - timedelta(seconds=CHAPTER_COVER_CLAIM_TTL_SECONDS)
    fresh_processing = db.execute(
        select(ChapterCoverIntent.id)
        .where(ChapterCoverIntent.chapter_id == chapter.id)
        .where(ChapterCoverIntent.status == "processing")
        .where(ChapterCoverIntent.claimed_at.is_not(None))
        .where(ChapterCoverIntent.claimed_at >= cutoff)
        .limit(1)
    ).scalar_one_or_none()
    if fresh_processing:
        return None

    intent = ChapterCoverIntent(
        id=str(uuid.uuid4()),
        chapter_id=chapter.id,
        prompt_brief=_build_chapter_cover_brief(chapter),
        status="processing",
        claim_token=claim_token,
        claimed_at=now,
        attempt_count=1,
    )
    with transactional_sync(db):
        db.add(intent)
        db.flush()
    return intent


@shared_task(bind=True, max_retries=3, default_retry_delay=30, ignore_result=True)
def generate_chapter_cover(self, chapter_id: str):
    """
    为 chapter 生成封面。
    从 chapter_cover_intents 原子认领 intent，或创建新 intent 后生成，
    写入 assets 并绑定到 chapters.cover_asset_id。
    """
    t0 = time.perf_counter()
    logger.info(
        "event=chapter_cover_task_start chapter_id={} msg=章节封面生成任务开始",
        chapter_id,
    )
    lock_key = f"lock:chapter-images:{chapter_id}"
    lock_handle = acquire_redis_lock(
        lock_key, ttl_seconds=CHAPTER_COVER_LOCK_TTL_SECONDS
    )
    if lock_handle is None:
        logger.debug("generate_chapter_cover: chapter={}, reason=locked", chapter_id)
        return {"status": "locked"}

    claim_token = uuid.uuid4().hex
    intent = None
    try:
        with get_sync_db() as db:
            stmt = (
                select(Chapter)
                .where(Chapter.id == chapter_id)
                .options(
                    joinedload(Chapter.story_links).joinedload(ChapterStoryLink.story),
                )
            )
            chapter = db.execute(stmt).unique().scalar_one_or_none()
            if not chapter:
                logger.debug(
                    "generate_chapter_cover: chapter={}, reason=not_found", chapter_id
                )
                return {"status": "no_chapter"}

            if not chapter_eligible_for_cover_by_inline_body_image_count(chapter):
                logger.debug(
                    "generate_chapter_cover: chapter={}, reason=insufficient_inline_body_images",
                    chapter_id,
                )
                return {"status": "insufficient_inline_body_images"}
            if not chapter_has_story_links(chapter):
                logger.debug(
                    "generate_chapter_cover: chapter={}, reason=no_story_links",
                    chapter_id,
                )
                return {"status": "no_story_links"}

            if getattr(chapter, "cover_asset_id", None):
                logger.debug(
                    "generate_chapter_cover: chapter={}, reason=has_cover_asset",
                    chapter_id,
                )
                return {"status": "already_has_asset"}

            intent = _claim_chapter_cover_intent_sync(db, chapter, claim_token)
            if not intent:
                logger.debug(
                    "generate_chapter_cover: chapter={}, reason=no_claimable_intent",
                    chapter_id,
                )
                return {"status": "no_intent"}

            chapter_title = chapter.title or ""
            chapter_category = chapter.category or ""
            chapter_user_id = chapter.user_id

        try:
            image_generator = get_image_generator()
            storage = TencentCosStorageService.from_env()
            from app.features.memoir.memoir_images.settings import MemoirImageSettings

            settings = MemoirImageSettings.from_env()
            orch = get_image_prompt_orchestrator()
            prompt_out = orch.build_cover_prompt(
                chapter_title=chapter_title,
                chapter_category=chapter_category,
                context_excerpt=intent.prompt_brief or "",
            )
            prompt_final = prompt_out["prompt"]
            style_for_image = prompt_out.get("style") or settings.default_style
            size_for_image = prompt_out.get("size") or settings.default_size

            result = image_generator.generate(
                prompt_final,
                size_for_image,
                style_for_image,
            )
            if result.status != TaskStatus.COMPLETED or not result.image_url:
                raise RuntimeError(result.error or "Image generation failed")

            image_bytes = _normalize_image_bytes(
                image_generator.download_image(result.image_url)
            )
            cos_key = _build_cover_cos_key(chapter_user_id, chapter_id, prompt_final)
            url = storage.upload_bytes(image_bytes, cos_key, "image/png")

            asset_id = str(uuid.uuid4())
            with get_sync_db() as db:
                intent_db = db.get(ChapterCoverIntent, intent.id)
                if (
                    not intent_db
                    or (intent_db.status or "").strip() != "processing"
                    or (intent_db.claim_token or "").strip() != claim_token
                ):
                    logger.debug(
                        "generate_chapter_cover: skip persist intent={} status={} claim={}",
                        intent.id,
                        getattr(intent_db, "status", None),
                        getattr(intent_db, "claim_token", None),
                    )
                    return {"status": "superseded_or_cancelled"}

                with transactional_sync(db):
                    asset = Asset(
                        id=asset_id,
                        asset_type="chapter_cover",
                        storage_key=cos_key,
                        url=url,
                        provider=settings.provider,
                        style_profile=style_for_image,
                        prompt_final=prompt_final,
                        status="completed",
                    )
                    db.add(asset)
                    db.flush()

                    intent_db.asset_id = asset_id
                    intent_db.status = "completed"
                    intent_db.claim_token = None
                    intent_db.claimed_at = None
                    intent_db.error = None
                    intent_db.updated_at = datetime.now(timezone.utc)

                    chapter_db = db.get(Chapter, chapter_id)
                    chapter_db.cover_asset_id = asset_id

            ms = (time.perf_counter() - t0) * 1000
            logger.info(
                "event=chapter_cover_task_done chapter_id={} asset_id={} duration_ms={:.1f} "
                "msg=章节封面生成完成",
                chapter_id,
                asset_id,
                ms,
            )
            logger.debug(
                "generate_chapter_cover: chapter={} asset={} url={} cos_key={} prompt_final={}",
                chapter_id,
                asset_id,
                url,
                cos_key,
                prompt_final,
            )
            return {"status": "success", "asset_id": asset_id}
        except Exception as exc:
            if intent is not None:
                with get_sync_db() as db:
                    with transactional_sync(db):
                        intent_db = db.get(ChapterCoverIntent, intent.id)
                        if (
                            intent_db
                            and (intent_db.claim_token or "").strip() == claim_token
                        ):
                            intent_db.status = "failed"
                            intent_db.claim_token = None
                            intent_db.claimed_at = None
                            intent_db.error = str(exc)
                            intent_db.updated_at = datetime.now(timezone.utc)
            ms = (time.perf_counter() - t0) * 1000
            logger.warning(
                "event=chapter_cover_task_failed chapter_id={} duration_ms={:.1f} error={} "
                "msg=章节封面生成失败",
                chapter_id,
                ms,
                exc,
            )
            raise self.retry(exc=exc) from exc
    finally:
        release_redis_lock(lock_handle)