video-gen/video-gen-api/app/services/generation_history_meta_service.py

from __future__ import annotations

from typing import Any, TypedDict

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.enums.generation_history import (
    GenerationHistorySourceEnum,
    get_generation_history_source_label,
    is_generation_history_module_source,
)
from app.models.module_generation_project import ModuleGenerationProject
from app.models.module_generation_step import ModuleGenerationStep
from app.models.shot_replicate_segment import ShotReplicateSegment


class GenerationHistoryMeta(TypedDict):
    """素材云历史项模块上下文回填字段。"""

    history_source: str | None
    history_source_label: str | None
    module_project_id: str | None
    module_project_title: str | None
    module_step_id: str | None
    module_step_code: str | None
    hot_opening_project_id: str | None
    hot_opening_project_title: str | None
    shot_replicate_project_id: str | None
    shot_replicate_project_title: str | None
    shot_task_set_id: str | None
    shot_segment_id: str | None
    shot_segment_index: int | None
    shot_segment_label: str | None


class _StepLinkInfo(TypedDict):
    module_project_id: str | None
    module_step_id: str | None
    module_step_code: str | None
    module: str | None


class _ProjectInfo(TypedDict):
    module_project_id: str
    module_project_title: str | None
    module: str | None


class _ShotSegmentInfo(TypedDict):
    shot_task_set_id: str | None
    shot_segment_id: str | None
    shot_segment_index: int | None
    shot_segment_label: str | None


def _unique(values: list[str] | tuple[str, ...]) -> list[str]:
    return list(dict.fromkeys(str(value) for value in values if value))


def _segment_label(segment_index: int | None) -> str | None:
    if segment_index is None:
        return None
    return f"片段{segment_index}"


def build_empty_history_meta(source: GenerationHistorySourceEnum) -> GenerationHistoryMeta:
    """构造统一历史字段，非对应模块字段保持 null。"""

    return {
        "history_source": source.value,
        "history_source_label": get_generation_history_source_label(source),
        "module_project_id": None,
        "module_project_title": None,
        "module_step_id": None,
        "module_step_code": None,
        "hot_opening_project_id": None,
        "hot_opening_project_title": None,
        "shot_replicate_project_id": None,
        "shot_replicate_project_title": None,
        "shot_task_set_id": None,
        "shot_segment_id": None,
        "shot_segment_index": None,
        "shot_segment_label": None,
    }


async def _load_step_link_map(
    db: AsyncSession,
    *,
    chat_task_ids: list[str],
    source: GenerationHistorySourceEnum,
) -> dict[str, _StepLinkInfo]:
    ids = _unique(chat_task_ids)
    if not ids:
        return {}

    stmt = (
        select(
            ModuleGenerationStep.chat_task_id.label("chat_task_id"),
            ModuleGenerationStep.id.label("module_step_id"),
            ModuleGenerationStep.project_id.label("module_project_id"),
            ModuleGenerationStep.step_code.label("module_step_code"),
            ModuleGenerationStep.module.label("module"),
            ModuleGenerationStep.is_current.label("is_current"),
            ModuleGenerationStep.updated_at.label("updated_at"),
        )
        .where(
            ModuleGenerationStep.deleted_at.is_(None),
            ModuleGenerationStep.chat_task_id.in_(ids),
            ModuleGenerationStep.module == source.value,
        )
        .order_by(
            ModuleGenerationStep.chat_task_id.asc(),
            ModuleGenerationStep.is_current.desc(),
            ModuleGenerationStep.updated_at.desc(),
        )
    )

    rows = (await db.execute(stmt)).mappings().all()
    link_map: dict[str, _StepLinkInfo] = {}
    for row in rows:
        chat_task_id = row["chat_task_id"]
        if not chat_task_id or chat_task_id in link_map:
            continue
        link_map[chat_task_id] = {
            "module_project_id": row["module_project_id"],
            "module_step_id": row["module_step_id"],
            "module_step_code": row["module_step_code"],
            "module": row["module"],
        }
    return link_map


async def _load_project_map(
    db: AsyncSession,
    *,
    module_project_ids: list[str],
) -> dict[str, _ProjectInfo]:
    ids = _unique(module_project_ids)
    if not ids:
        return {}

    stmt = (
        select(
            ModuleGenerationProject.id.label("module_project_id"),
            ModuleGenerationProject.title.label("module_project_title"),
            ModuleGenerationProject.module.label("module"),
        )
        .where(
            ModuleGenerationProject.deleted_at.is_(None),
            ModuleGenerationProject.id.in_(ids),
        )
    )

    return {
        row["module_project_id"]: {
            "module_project_id": row["module_project_id"],
            "module_project_title": row["module_project_title"],
            "module": row["module"],
        }
        for row in (await db.execute(stmt)).mappings().all()
        if row["module_project_id"]
    }


async def _load_shot_segment_map(
    db: AsyncSession,
    *,
    module_project_ids: list[str],
) -> dict[str, _ShotSegmentInfo]:
    ids = _unique(module_project_ids)
    if not ids:
        return {}

    stmt = (
        select(
            ShotReplicateSegment.module_project_id.label("module_project_id"),
            ShotReplicateSegment.id.label("shot_segment_id"),
            ShotReplicateSegment.task_set_id.label("shot_task_set_id"),
            ShotReplicateSegment.segment_index.label("shot_segment_index"),
            ShotReplicateSegment.updated_at.label("updated_at"),
        )
        .where(
            ShotReplicateSegment.deleted_at.is_(None),
            ShotReplicateSegment.module_project_id.in_(ids),
        )
        .order_by(
            ShotReplicateSegment.module_project_id.asc(),
            ShotReplicateSegment.updated_at.desc(),
        )
    )

    segment_map: dict[str, _ShotSegmentInfo] = {}
    rows = (await db.execute(stmt)).mappings().all()
    for row in rows:
        module_project_id = row["module_project_id"]
        if not module_project_id or module_project_id in segment_map:
            continue
        segment_index = row["shot_segment_index"]
        segment_map[module_project_id] = {
            "shot_task_set_id": row["shot_task_set_id"],
            "shot_segment_id": row["shot_segment_id"],
            "shot_segment_index": segment_index,
            "shot_segment_label": _segment_label(segment_index),
        }
    return segment_map


def _merge_meta(
    *,
    source: GenerationHistorySourceEnum,
    step_info: _StepLinkInfo | None,
    project_info: _ProjectInfo | None,
    shot_info: _ShotSegmentInfo | None,
) -> GenerationHistoryMeta:
    meta = build_empty_history_meta(source)

    module_project_id = step_info["module_project_id"] if step_info else None
    module_step_id = step_info["module_step_id"] if step_info else None
    module_step_code = step_info["module_step_code"] if step_info else None
    module_project_title = project_info["module_project_title"] if project_info else None

    meta["module_project_id"] = module_project_id
    meta["module_project_title"] = module_project_title
    meta["module_step_id"] = module_step_id
    meta["module_step_code"] = module_step_code

    if source == GenerationHistorySourceEnum.HOT_OPENING_REPLICATE:
        meta["hot_opening_project_id"] = module_project_id
        meta["hot_opening_project_title"] = module_project_title
    elif source == GenerationHistorySourceEnum.SHOT_REPLICATE:
        meta["shot_replicate_project_id"] = module_project_id
        meta["shot_replicate_project_title"] = module_project_title
        if shot_info:
            meta["shot_task_set_id"] = shot_info["shot_task_set_id"]
            meta["shot_segment_id"] = shot_info["shot_segment_id"]
            meta["shot_segment_index"] = shot_info["shot_segment_index"]
            meta["shot_segment_label"] = shot_info["shot_segment_label"]

    return meta


async def batch_load_generation_history_meta_map(
    db: AsyncSession,
    *,
    source: GenerationHistorySourceEnum,
    chat_task_ids: list[str],
) -> dict[str, GenerationHistoryMeta]:
    """批量回填历史列表模块上下文，避免逐条链式查询。"""

    ids = _unique(chat_task_ids)
    if not ids:
        return {}

    if not is_generation_history_module_source(source):
        return {chat_task_id: build_empty_history_meta(source) for chat_task_id in ids}

    step_link_map = await _load_step_link_map(db, chat_task_ids=ids, source=source)
    module_project_ids = [
        step_info["module_project_id"]
        for step_info in step_link_map.values()
        if step_info.get("module_project_id")
    ]
    project_map = await _load_project_map(db, module_project_ids=module_project_ids)

    shot_segment_map: dict[str, _ShotSegmentInfo] = {}
    if source == GenerationHistorySourceEnum.SHOT_REPLICATE:
        shot_segment_map = await _load_shot_segment_map(db, module_project_ids=module_project_ids)

    meta_map: dict[str, GenerationHistoryMeta] = {}
    for chat_task_id in ids:
        step_info = step_link_map.get(chat_task_id)
        module_project_id = step_info.get("module_project_id") if step_info else None
        project_info = project_map.get(module_project_id) if module_project_id else None
        shot_info = shot_segment_map.get(module_project_id) if module_project_id else None
        meta_map[chat_task_id] = _merge_meta(
            source=source,
            step_info=step_info,
            project_info=project_info,
            shot_info=shot_info,
        )
    return meta_map