SillyTavern_replica/backend/services/fiction_orchestrator_service.py

"""
爽文阅读流水线编排 — 新版 ensure 滚动补齐：卷纲 → 事件链 → 章纲 → 章节。
"""
from __future__ import annotations

import asyncio
import logging
from typing import Dict, List, Optional

from models.fiction_models import (
    FictionPipelineSettings,
    FictionPipelineTickResult,
    FictionRunState,
    FictionStartReadingResult,
)
from services.fiction_chapter_service import (
    find_next_unwritten_chapter,
    has_written_chapters,
    run_chapter,
)
from services.fiction_metadata_service import fiction_metadata_service
from services.fiction_planning_service import (
    ensure_chapter_plan,
    ensure_event_chain,
    ensure_volume,
)
from services.fiction_service import fiction_service

logger = logging.getLogger(__name__)

_active_tasks: Dict[str, asyncio.Task] = {}
_lock = asyncio.Lock()


def _get_pipeline_settings(book_id: str) -> FictionPipelineSettings:
    settings = fiction_service.get_book_settings(book_id)
    return settings.pipeline or FictionPipelineSettings()


def _needs_volume(book_id: str) -> bool:
    metadata = fiction_metadata_service.get_metadata(book_id)
    return not metadata.volumes


def _needs_event_chain(book_id: str) -> bool:
    metadata = fiction_metadata_service.get_metadata(book_id)
    if not metadata.volumes:
        return False
    for volume in metadata.volumes:
        if not metadata.eventChains.get(volume.id):
            return True
    return False


def _needs_chapter_plan(book_id: str) -> bool:
    metadata = fiction_metadata_service.get_metadata(book_id)
    if not metadata.volumes:
        return False
    for volume in metadata.volumes:
        events = metadata.eventChains.get(volume.id, [])
        if not events:
            return False
        for event in events:
            if not metadata.chapterPlans.get(event.id):
                return True
    return False


def _needs_chapter(book_id: str) -> bool:
    if _needs_volume(book_id) or _needs_event_chain(book_id) or _needs_chapter_plan(book_id):
        return False
    if has_written_chapters(book_id):
        return False
    metadata = fiction_metadata_service.get_metadata(book_id)
    return find_next_unwritten_chapter(book_id, metadata) is not None


def _pipeline_complete(book_id: str) -> bool:
    return (
        not _needs_volume(book_id)
        and not _needs_event_chain(book_id)
        and not _needs_chapter_plan(book_id)
        and not _needs_chapter(book_id)
    )


def get_pending_stages(book_id: str) -> List[str]:
    """返回需手动触发的阶段 id 列表（auto 关闭且仍有工作，或上次失败需重试）。"""
    pipeline = _get_pipeline_settings(book_id)
    pending: List[str] = []
    if _needs_volume(book_id) and not pipeline.autoCoarse:
        pending.append("volume")
    if _needs_event_chain(book_id) and not pipeline.autoEventPlan:
        pending.append("event_chain")
    if _needs_chapter_plan(book_id) and not pipeline.autoEventPlan:
        pending.append("chapter_plan")
    if _needs_chapter(book_id) and not pipeline.autoChapter:
        pending.append("chapter")

    run = fiction_metadata_service.get_run(book_id)
    if run.status == "error" and run.pipelineStage:
        retry_map = {
            "volume": "volume",
            "coarse": "volume",
            "event_chain": "event_chain",
            "event_plan": "chapter_plan",
            "chapter_plan": "chapter_plan",
            "chapter": "chapter",
        }
        failed = retry_map.get(run.pipelineStage)
        if failed and failed not in pending:
            if failed == "volume" and _needs_volume(book_id):
                pending.insert(0, failed)
            elif failed == "event_chain" and _needs_event_chain(book_id):
                pending.insert(0, failed)
            elif failed == "chapter_plan" and _needs_chapter_plan(book_id):
                pending.insert(0, failed)
            elif failed == "chapter" and _needs_chapter(book_id):
                pending.insert(0, failed)
    return pending


def _next_auto_stage(book_id: str) -> Optional[str]:
    pipeline = _get_pipeline_settings(book_id)
    if _needs_volume(book_id):
        return "volume" if pipeline.autoCoarse else None
    if _needs_event_chain(book_id):
        return "event_chain" if pipeline.autoEventPlan else None
    if _needs_chapter_plan(book_id):
        return "chapter_plan" if pipeline.autoEventPlan else None
    if _needs_chapter(book_id):
        return "chapter" if pipeline.autoChapter else None
    return None


async def _run_pipeline(
    book_id: str,
    *,
    profile_id: Optional[str] = None,
    api_config: Optional[Dict[str, str]] = None,
) -> None:
    pipeline = _get_pipeline_settings(book_id)
    try:
        if _needs_volume(book_id):
            if not pipeline.autoCoarse:
                return
            await ensure_volume(
                book_id, profile_id=profile_id, api_config=api_config
            )

        if _needs_event_chain(book_id):
            if not pipeline.autoEventPlan:
                return
            await ensure_event_chain(
                book_id, profile_id=profile_id, api_config=api_config
            )

        if _needs_chapter_plan(book_id):
            if not pipeline.autoEventPlan:
                return
            await ensure_chapter_plan(
                book_id, profile_id=profile_id, api_config=api_config
            )

        if _needs_chapter(book_id):
            if not pipeline.autoChapter:
                return
            await run_chapter(
                book_id, profile_id=profile_id, api_config=api_config
            )

        if _pipeline_complete(book_id):
            fiction_metadata_service.set_pipeline_stage(
                book_id, status="idle", pipeline_stage="ready"
            )
    except Exception:
        logger.exception("Fiction pipeline failed for book %s", book_id)
    finally:
        async with _lock:
            _active_tasks.pop(book_id, None)


async def _task_is_active(book_id: str) -> bool:
    async with _lock:
        task = _active_tasks.get(book_id)
        return task is not None and not task.done()


async def _start_pipeline_task(
    book_id: str,
    *,
    profile_id: Optional[str] = None,
    api_config: Optional[Dict[str, str]] = None,
) -> FictionPipelineTickResult:
    run = fiction_metadata_service.get_run(book_id)
    pending = get_pending_stages(book_id)

    if run.status == "running":
        if await _task_is_active(book_id):
            return FictionPipelineTickResult(run=run, started=False, pendingStages=pending)
        logger.warning("Stale running pipeline for book %s, resetting to idle", book_id)
        run = fiction_metadata_service.set_pipeline_stage(
            book_id, status="idle", pipeline_stage=run.pipelineStage
        )

    if run.status == "error":
        if _pipeline_complete(book_id):
            run = fiction_metadata_service.set_pipeline_stage(
                book_id, status="idle", pipeline_stage="ready"
            )
            return FictionPipelineTickResult(
                run=run, started=False, pendingStages=pending
            )
        run = fiction_metadata_service.clear_pipeline_error(book_id)

    if _pipeline_complete(book_id):
        if run.pipelineStage != "ready":
            run = fiction_metadata_service.set_pipeline_stage(
                book_id, status="idle", pipeline_stage="ready"
            )
        return FictionPipelineTickResult(run=run, started=False, pendingStages=pending)

    next_stage = _next_auto_stage(book_id)
    if not next_stage:
        if run.pipelineStage not in (
            None,
            "ready",
            "volume_done",
            "event_chain_done",
            "chapter_plan_done",
            "chapter_done",
        ):
            run = fiction_metadata_service.set_pipeline_stage(
                book_id, status="idle", pipeline_stage="ready"
            )
        return FictionPipelineTickResult(
            run=run, started=False, pendingStages=pending
        )

    async with _lock:
        existing = _active_tasks.get(book_id)
        if existing and not existing.done():
            run = fiction_metadata_service.get_run(book_id)
            return FictionPipelineTickResult(
                run=run, started=False, pendingStages=pending
            )

        fiction_metadata_service.set_pipeline_stage(
            book_id, status="running", pipeline_stage=next_stage
        )

        task = asyncio.create_task(
            _run_pipeline(
                book_id, profile_id=profile_id, api_config=api_config
            )
        )
        _active_tasks[book_id] = task

    run = fiction_metadata_service.get_run(book_id)
    pending = get_pending_stages(book_id)
    return FictionPipelineTickResult(run=run, started=True, pendingStages=pending)


async def tick_reading_pipeline(
    book_id: str,
    *,
    profile_id: Optional[str] = None,
    api_config: Optional[Dict[str, str]] = None,
) -> FictionPipelineTickResult:
    """检查 metadata + settings，按需启动下一自动阶段。"""
    return await _start_pipeline_task(
        book_id, profile_id=profile_id, api_config=api_config
    )


async def start_reading_pipeline(
    book_id: str,
    *,
    profile_id: Optional[str] = None,
    api_config: Optional[Dict[str, str]] = None,
) -> FictionStartReadingResult:
    """进入阅读时的流水线入口（兼容旧接口）。"""
    result = await tick_reading_pipeline(
        book_id, profile_id=profile_id, api_config=api_config
    )
    return FictionStartReadingResult(run=result.run, started=result.started)


def get_pipeline_run(book_id: str) -> FictionRunState:
    return fiction_metadata_service.get_run(book_id)