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

"""
Story 主插图生成 Celery 任务。

从 story_image_intents 原子 claim intent，生成图片，写入 assets，更新 intent。
不读取正文占位符。
"""

import hashlib
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 app.agents.image_prompt import get_image_prompt_orchestrator
from app.core.db import get_sync_db
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.asset_resolver import strip_asset_image_refs_from_markdown
from app.features.memoir.memoir_images.settings import MemoirImageSettings
from app.features.memoir.memoir_images.storage import TencentCosStorageService
from app.features.story.backfill import backfill_image_into_markdown
from app.features.story.models import Story, StoryImageIntent, StoryVersion
from app.ports.image_gen import TaskStatus

logger = get_logger(__name__)

STORY_IMAGE_LOCK_TTL_SECONDS = 1800
STORY_IMAGE_CLAIM_TTL_SECONDS = 1800


def _enqueue_chapter_effects_after_image_backfill(story_id: str) -> None:
    """主图回填后标记关联章节 dirty，并经统一 post-commit 入口派发章节物化与 compaction。"""
    try:
        with get_sync_db() as session:
            from app.features.memoir import repo as memoir_repo

            story = session.get(Story, story_id)
            if not story:
                return
            uid = str(story.user_id)
            memoir_repo.mark_chapters_dirty_for_story_sync(session, story_id)
            chapter_ids = memoir_repo.get_chapter_ids_linked_to_story_sync(
                session, story_id
            )
            session.commit()
            user_id = uid
    except Exception as exc:
        logger.warning(
            "mark_chapters_dirty_for_story_sync failed story={}: {}", story_id, exc
        )
        return
    from app.features.story.post_commit import enqueue_story_post_commit_effects

    enqueue_story_post_commit_effects(
        user_id=user_id,
        story_ids=set(),
        chapter_ids=set(chapter_ids),
        trigger_source="story_image_backfill",
        need_image=False,
        need_recompose=True,
        need_compaction=True,
    )


def _build_story_image_cos_key(
    user_id: str, story_id: str, intent_id: str, prompt: str
) -> str:
    short_hash = hashlib.sha1(prompt.encode("utf-8")).hexdigest()[:10]
    return f"stories/{user_id}/{story_id}/{intent_id}-{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 _story_image_claimable_clause(now: datetime):
    cutoff = now - timedelta(seconds=STORY_IMAGE_CLAIM_TTL_SECONDS)
    return or_(
        StoryImageIntent.status.in_(["pending", "failed"]),
        and_(
            StoryImageIntent.status == "processing",
            or_(
                StoryImageIntent.claimed_at.is_(None),
                StoryImageIntent.claimed_at < cutoff,
            ),
        ),
    )


def _claim_story_image_intent_sync(db, story_id: str, claim_token: str):
    now = datetime.now(timezone.utc)
    claimable = _story_image_claimable_clause(now)
    candidate_id = db.execute(
        select(StoryImageIntent.id)
        .where(StoryImageIntent.story_id == story_id)
        .where(StoryImageIntent.intent_role == "primary")
        .where(claimable)
        .order_by(
            StoryImageIntent.updated_at.desc(), StoryImageIntent.created_at.desc()
        )
        .limit(1)
    ).scalar_one_or_none()
    if not candidate_id:
        return None

    claimed = db.execute(
        update(StoryImageIntent)
        .where(StoryImageIntent.id == candidate_id)
        .where(_story_image_claimable_clause(now))
        .values(
            status="processing",
            claim_token=claim_token,
            claimed_at=now,
            updated_at=now,
            error=None,
            attempt_count=func.coalesce(StoryImageIntent.attempt_count, 0) + 1,
        )
    )
    if (claimed.rowcount or 0) != 1:
        db.rollback()
        return None

    row = (
        db.execute(
            select(StoryImageIntent, Story)
            .join(Story, StoryImageIntent.story_id == Story.id)
            .where(StoryImageIntent.id == candidate_id)
        )
        .unique()
        .first()
    )
    db.commit()
    return row


@shared_task(bind=True, max_retries=3, default_retry_delay=30)
def generate_story_image(self, story_id: str):
    """
    为 story 生成主插图。
    从 story_image_intents 原子认领 primary intent，生成后写入 assets 并更新 intent。
    """
    lock_key = f"lock:story-image:{story_id}"
    lock_handle = acquire_redis_lock(lock_key, ttl_seconds=STORY_IMAGE_LOCK_TTL_SECONDS)
    if lock_handle is None:
        logger.debug("generate_story_image: story={}, reason=locked", story_id)
        return {"status": "locked"}

    claim_token = uuid.uuid4().hex
    intent = None
    story = None
    try:
        with get_sync_db() as db:
            row = _claim_story_image_intent_sync(db, story_id, claim_token)
            if not row:
                logger.debug(
                    "generate_story_image: story={}, reason=no_claimable_intent",
                    story_id,
                )
                return {"status": "no_intent"}

            intent, story = row

        img_cfg = MemoirImageSettings.from_env()
        min_body = img_cfg.story_image_min_body_chars
        if min_body > 0:
            plain = strip_asset_image_refs_from_markdown(
                story.canonical_markdown or ""
            ).strip()
            if len(plain) < min_body:
                with get_sync_db() as db:
                    intent_db = db.get(StoryImageIntent, intent.id)
                    if intent_db and (intent_db.status or "").strip() == "processing":
                        intent_db.status = "skipped"
                        intent_db.error = f"body_below_min_chars:{len(plain)}"
                        intent_db.claim_token = None
                        intent_db.claimed_at = None
                        intent_db.updated_at = datetime.now(timezone.utc)
                        db.commit()
                logger.info(
                    "generate_story_image: skipped body too short story={} len={} min={}",
                    story_id,
                    len(plain),
                    min_body,
                )
                return {"status": "skipped_body_too_short"}

        image_generator = get_image_generator()
        storage = TencentCosStorageService.from_env()

        settings = img_cfg
        orch = get_image_prompt_orchestrator()
        intent_style = (getattr(intent, "style_profile", None) or "").strip() or None
        prompt_out = orch.build_story_primary_prompt(
            story_title=story.title or "",
            story_stage=story.stage,
            prompt_brief=intent.prompt_brief or "",
            style_profile=intent_style,
        )
        prompt_final = prompt_out["prompt"]
        style_for_image = (
            (prompt_out.get("style") or "").strip()
            or intent_style
            or settings.default_style
        )
        result = image_generator.generate(
            prompt_final,
            prompt_out.get("size") or settings.default_size,
            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_story_image_cos_key(
            story.user_id, story_id, intent.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(StoryImageIntent, 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_story_image: skip persist intent={} status={} claim={}",
                    intent.id,
                    getattr(intent_db, "status", None),
                    getattr(intent_db, "claim_token", None),
                )
                return {"status": "superseded_or_cancelled"}

            asset = Asset(
                id=asset_id,
                asset_type="story_image",
                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()

            story_db = db.get(Story, story_id)
            target_vid = intent_db.story_version_id or story_db.current_version_id
            current_vid = story_db.current_version_id

            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)
            db.flush()

            # 仅当 intent 仍指向当前版本时回填正文，避免慢任务/重试把图插到新版本上
            if not target_vid or target_vid != current_vid:
                db.commit()
                logger.debug(
                    "generate_story_image: stale intent skip backfill story={} "
                    "intent_ver={} current={} url={} asset={}",
                    story_id,
                    target_vid,
                    current_vid,
                    url,
                    asset_id,
                )
                return {"status": "success_stale", "asset_id": asset_id}

            ver = db.get(StoryVersion, target_vid)
            if not ver:
                db.commit()
                return {"status": "success_no_snapshot", "asset_id": asset_id}

            base_md = strip_asset_image_refs_from_markdown(ver.markdown_snapshot or "")
            alt_text = (getattr(intent_db, "prompt_brief", None) or "").strip()
            if not alt_text:
                alt_text = (getattr(intent_db, "caption", None) or "").strip()
            backfilled_md = backfill_image_into_markdown(
                base_md,
                asset_id=asset_id,
                image_alt=alt_text or "主插图",
            )
            max_stmt = select(func.max(StoryVersion.version_no)).where(
                StoryVersion.story_id == story_id
            )
            max_no = db.execute(max_stmt).scalar()
            version_no = (max_no or 0) + 1
            new_ver = StoryVersion(
                id=str(uuid.uuid4()),
                story_id=story_id,
                version_no=version_no,
                markdown_snapshot=backfilled_md,
                change_summary="主插图回填",
                actor_type="system",
                source_type="image_backfill",
                parent_version_id=story_db.current_version_id,
            )
            db.add(new_ver)
            db.flush()
            story_db.current_version_id = new_ver.id
            story_db.canonical_markdown = backfilled_md

            db.commit()

        _enqueue_chapter_effects_after_image_backfill(story_id)

        logger.info(
            "generate_story_image: story={}, asset={}",
            story_id,
            asset_id,
        )
        logger.debug(
            "generate_story_image: story={} asset={} url={} cos_key={} prompt_final={}",
            story_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:
                intent_db = db.get(StoryImageIntent, intent.id)
                if (
                    intent_db
                    and (intent_db.status or "").strip() != "completed"
                    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)
                    db.commit()
        logger.warning("generate_story_image failed: story={}, error={}", story_id, exc)
        raise self.retry(exc=exc) from exc
    finally:
        release_redis_lock(lock_handle)