SillyTavern_replica/backend/services/fiction_metadata_service.py

"""
爽文 metadata.json / run.json 读写服务。
"""
from __future__ import annotations

import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional

from models.fiction_models import FictionBookMetadata, FictionRunState
from services.fiction_service import fiction_service

logger = logging.getLogger(__name__)

_STAGE_MESSAGES: Dict[tuple, tuple] = {
    ("running", "coarse"): ("coarse_generating", "正在生成粗纲…"),
    ("running", "event_plan"): ("event_plan_generating", "正在生成事件纲要…"),
    ("running", "chapter"): ("chapter_generating", "正在撰写正文…"),
    ("error", "coarse"): ("error", "粗纲生成失败"),
    ("error", "event_plan"): ("error", "事件纲要生成失败"),
    ("error", "chapter"): ("error", "章节写作失败"),
}


def _resolve_stage_message(
    status: str,
    pipeline_stage: Optional[str],
    override_message: Optional[str] = None,
) -> tuple:
    if override_message is not None:
        key = (status, pipeline_stage or "")
        stage = _STAGE_MESSAGES.get(key, (status, override_message))[0]
        if status == "error":
            stage = "error"
        elif status == "running" and pipeline_stage == "coarse":
            stage = "coarse_generating"
        elif status == "running" and pipeline_stage == "event_plan":
            stage = "event_plan_generating"
        elif status == "running" and pipeline_stage == "chapter":
            stage = "chapter_generating"
        elif status == "idle":
            stage = "idle"
        return stage, override_message
    matched = _STAGE_MESSAGES.get((status, pipeline_stage or ""))
    if matched:
        return matched
    if status == "idle":
        return "idle", None
    if status == "error":
        return "error", "生成失败"
    return status, None


def _read_json(path: Path) -> Any:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def _write_json(path: Path, data: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


class FictionMetadataService:
    def _metadata_path(self, book_id: str) -> Path:
        return fiction_service._metadata_path(book_id)

    def _run_path(self, book_id: str) -> Path:
        return fiction_service._run_path(book_id)

    def _ensure_book(self, book_id: str) -> None:
        if not fiction_service._meta_path(book_id).exists():
            raise FileNotFoundError(f"Book not found: {book_id}")

    def get_metadata(self, book_id: str) -> FictionBookMetadata:
        self._ensure_book(book_id)
        path = self._metadata_path(book_id)
        if not path.exists():
            raise FileNotFoundError(f"metadata.json not found for book: {book_id}")
        return FictionBookMetadata(**_read_json(path))

    def save_metadata(self, book_id: str, metadata: FictionBookMetadata) -> FictionBookMetadata:
        self._ensure_book(book_id)
        _write_json(self._metadata_path(book_id), metadata.model_dump())
        self._touch_book_meta(book_id)
        return metadata

    def update_metadata(self, book_id: str, patch: Dict[str, Any]) -> FictionBookMetadata:
        current = self.get_metadata(book_id)
        data = current.model_dump()
        for key, value in patch.items():
            data[key] = value
        updated = FictionBookMetadata(**data)
        return self.save_metadata(book_id, updated)

    def get_run(self, book_id: str) -> FictionRunState:
        self._ensure_book(book_id)
        path = self._run_path(book_id)
        if not path.exists():
            return FictionRunState()
        return FictionRunState(**_read_json(path))

    def save_run(self, book_id: str, run: FictionRunState) -> FictionRunState:
        self._ensure_book(book_id)
        _write_json(self._run_path(book_id), run.model_dump())
        return run

    def set_pipeline_stage(
        self,
        book_id: str,
        *,
        status: str,
        pipeline_stage: Optional[str] = None,
        message: Optional[str] = None,
    ) -> FictionRunState:
        run = self.get_run(book_id)
        run.status = status
        run.pipelineStage = pipeline_stage
        run.stage, run.message = _resolve_stage_message(status, pipeline_stage, message)
        run.updatedAt = datetime.now().isoformat()
        return self.save_run(book_id, run)

    def clear_pipeline_error(self, book_id: str) -> FictionRunState:
        """清除 error 状态，保留 pipelineStage 供重试参考。"""
        run = self.get_run(book_id)
        if run.status != "error":
            return run
        run.status = "idle"
        run.stage = "idle"
        run.message = None
        run.updatedAt = datetime.now().isoformat()
        return self.save_run(book_id, run)

    def set_pipeline_progress(
        self, book_id: str, *, done: int, total: int
    ) -> FictionRunState:
        run = self.get_run(book_id)
        run.progress = {"done": done, "total": total}
        run.updatedAt = datetime.now().isoformat()
        return self.save_run(book_id, run)

    def clear_pipeline_progress(self, book_id: str) -> FictionRunState:
        run = self.get_run(book_id)
        run.progress = None
        run.updatedAt = datetime.now().isoformat()
        return self.save_run(book_id, run)

    def update_progress(
        self,
        book_id: str,
        *,
        current_chapter_seq: Optional[int] = None,
        char_offset: Optional[int] = None,
    ):
        metadata = self.get_metadata(book_id)
        progress = metadata.progress
        if current_chapter_seq is not None:
            progress.currentChapterSeq = current_chapter_seq
        if char_offset is not None:
            progress.charOffset = char_offset
        metadata.progress = progress
        return self.save_metadata(book_id, metadata)


fiction_metadata_service = FictionMetadataService()