SillyTavern_replica/backend/services/studio_project_service.py

"""
Load/save Studio projects and skill templates from data/agent/.
"""
from __future__ import annotations

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

from core.config import settings
from models.studio_models import (
    CreateStudioProjectRequest,
    PipelineDefinition,
    StudioProject,
    StudioProjectMeta,
    StudioProjectSummary,
    WorkflowTemplateSummary,
    WorkflowVariablesResponse,
    WorkflowVariableDef,
)

logger = logging.getLogger(__name__)

DEFAULT_TEMPLATE_ID = "builtin.studio.example"

POSITION_STRING_MAP = {
    "after_char": 0,
    "before_char": 1,
    "before_example": 2,
    "after_example": 3,
    "system": 4,
    "as_system": 5,
    "depth": 6,
    "macro": 7,
}

ACTIVATION_LEGACY_MAP = {
    "normal": "permanent",
    "constant": "permanent",
    "selective": "keyword",
}


def _normalize_position(value: Any) -> int:
    if value is None:
        return 1
    if isinstance(value, int):
        return value
    if isinstance(value, str):
        if value.isdigit():
            return int(value)
        return POSITION_STRING_MAP.get(value, 1)
    try:
        return int(value)
    except (TypeError, ValueError):
        return 1


def _normalize_activation(value: Any) -> str:
    if not value:
        return "permanent"
    text = str(value)
    return ACTIVATION_LEGACY_MAP.get(text, text)


def _migrate_scoring(scoring: Dict[str, Any]) -> Dict[str, Any]:
    if not scoring:
        return {"enabled": True, "dimensions": []}
    if scoring.get("dimensions"):
        return scoring
    rubric = scoring.get("rubric")
    if rubric:
        scoring = {**scoring}
        scoring["dimensions"] = [
            {
                "id": "default",
                "name": "综合质量",
                "criteria": rubric,
            }
        ]
        scoring.pop("rubric", None)
    return scoring


def _normalize_node(node: Dict[str, Any]) -> Dict[str, Any]:
    node = dict(node)
    config = dict(node.get("config") or {})
    insertion = dict(config.get("insertion") or {})
    if insertion:
        insertion["position"] = _normalize_position(insertion.get("position"))
        insertion["activationType"] = _normalize_activation(
            insertion.get("activationType")
        )
        config["insertion"] = insertion
    if "scoring" in config:
        config["scoring"] = _migrate_scoring(dict(config.get("scoring") or {}))
    node["config"] = config
    return node


def _parse_node_ref(ref: str, node_ids: set[str]) -> Optional[str]:
    if not ref or not ref.endswith(".output"):
        return None
    node_id = ref[: -len(".output")]
    return node_id if node_id in node_ids else None


def _build_node_dependency_edges(pipeline: Dict[str, Any]) -> List[tuple[str, str]]:
    nodes = pipeline.get("nodes") or []
    node_ids = {n["id"] for n in nodes if n.get("id")}
    edges: List[tuple[str, str]] = []
    seen: set[tuple[str, str]] = set()
    for node in nodes:
        to_id = node.get("id")
        if not to_id:
            continue
        for inp in node.get("inputs") or []:
            src = _parse_node_ref(inp.get("ref", ""), node_ids)
            if not src or src == to_id:
                continue
            pair = (src, to_id)
            if pair in seen:
                continue
            seen.add(pair)
            edges.append(pair)
    return edges


def _detect_reference_cycles(pipeline: Dict[str, Any]) -> List[List[str]]:
    nodes = pipeline.get("nodes") or []
    node_ids = [n["id"] for n in nodes if n.get("id")]
    adj: Dict[str, List[str]] = {nid: [] for nid in node_ids}
    for src, dst in _build_node_dependency_edges(pipeline):
        adj[src].append(dst)

    cycles: List[List[str]] = []
    visited: set[str] = set()
    stack: set[str] = set()
    path: List[str] = []

    def dfs(node_id: str) -> None:
        visited.add(node_id)
        stack.add(node_id)
        path.append(node_id)
        for nxt in adj.get(node_id, []):
            if nxt not in visited:
                dfs(nxt)
            elif nxt in stack:
                start = path.index(nxt)
                if start >= 0:
                    cycles.append(path[start:] + [nxt])
        path.pop()
        stack.discard(node_id)

    for nid in node_ids:
        if nid not in visited:
            dfs(nid)
    return cycles


def _validate_pipeline_refs(pipeline: Dict[str, Any]) -> None:
    cycles = _detect_reference_cycles(pipeline)
    if not cycles:
        return
    nodes = {n["id"]: n.get("displayName", n["id"]) for n in pipeline.get("nodes") or []}
    first = cycles[0]
    chain = " → ".join(nodes.get(nid, nid) for nid in first)
    raise ValueError(f"流水线存在循环引用：{chain}")


def _normalize_pipeline_dict(pipeline: Dict[str, Any]) -> Dict[str, Any]:
    pipeline = dict(pipeline)
    nodes = pipeline.get("nodes") or []
    pipeline["nodes"] = [_normalize_node(n) for n in nodes]
    return pipeline


def _read_json(path: Path) -> Any:
    with path.open("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 path.open("w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def _slugify(name: str) -> str:
    slug = re.sub(r"[^\w\u4e00-\u9fff-]+", "-", name.strip(), flags=re.UNICODE)
    slug = re.sub(r"-+", "-", slug).strip("-").lower()
    return slug or "project"


class StudioProjectService:
    @property
    def projects_root(self) -> Path:
        return settings.AGENT_STUDIO_PROJECTS_PATH

    @property
    def templates_root(self) -> Path:
        return settings.AGENT_TEMPLATES_PATH

    def _project_dir(self, project_id: str) -> Path:
        return self.projects_root / project_id

    def _meta_path(self, project_id: str) -> Path:
        return self._project_dir(project_id) / "meta.json"

    def _pipeline_path(self, project_id: str) -> Path:
        return self._project_dir(project_id) / "pipeline.json"

    def list_projects(self) -> List[StudioProjectSummary]:
        root = self.projects_root
        if not root.exists():
            return []
        summaries: List[StudioProjectSummary] = []
        for child in sorted(root.iterdir()):
            if not child.is_dir():
                continue
            meta_path = child / "meta.json"
            if not meta_path.exists():
                continue
            meta = _read_json(meta_path)
            summaries.append(
                StudioProjectSummary(
                    id=meta.get("id", child.name),
                    name=meta.get("name", child.name),
                    description=meta.get("description", ""),
                    updatedAt=meta.get("updatedAt", ""),
                )
            )
        return summaries

    def get_project(self, project_id: str) -> StudioProject:
        meta_path = self._meta_path(project_id)
        pipeline_path = self._pipeline_path(project_id)
        if not meta_path.exists() or not pipeline_path.exists():
            raise FileNotFoundError(f"Studio project not found: {project_id}")
        meta = _read_json(meta_path)
        pipeline = _normalize_pipeline_dict(_read_json(pipeline_path))
        return StudioProject(
            meta=StudioProjectMeta(**meta),
            pipeline=PipelineDefinition(**pipeline),
        )

    def update_project_bindings(
        self,
        project_id: str,
        character_id: str,
        worldbook_id: str,
    ) -> StudioProject:
        meta_path = self._meta_path(project_id)
        if not meta_path.exists():
            raise FileNotFoundError(f"Studio project not found: {project_id}")
        meta = _read_json(meta_path)
        meta["characterId"] = character_id
        meta["worldbookId"] = worldbook_id
        meta["updatedAt"] = datetime.now().isoformat()
        _write_json(meta_path, meta)
        return self.get_project(project_id)

    def update_project_meta(
        self,
        project_id: str,
        *,
        name: Optional[str] = None,
        description: Optional[str] = None,
    ) -> StudioProject:
        meta_path = self._meta_path(project_id)
        if not meta_path.exists():
            raise FileNotFoundError(f"Studio project not found: {project_id}")
        meta = _read_json(meta_path)
        if name is not None:
            meta["name"] = name.strip()
        if description is not None:
            meta["description"] = description
        meta["updatedAt"] = datetime.now().isoformat()
        _write_json(meta_path, meta)
        return self.get_project(project_id)

    def save_pipeline(self, project_id: str, pipeline: PipelineDefinition) -> StudioProject:
        meta_path = self._meta_path(project_id)
        if not meta_path.exists():
            raise FileNotFoundError(f"Studio project not found: {project_id}")
        normalized = _normalize_pipeline_dict(pipeline.model_dump(exclude_none=True))
        _validate_pipeline_refs(normalized)
        meta = _read_json(meta_path)
        now = datetime.now().isoformat()
        meta["updatedAt"] = now
        _write_json(meta_path, meta)
        _write_json(self._pipeline_path(project_id), normalized)
        return self.get_project(project_id)

    def list_workflow_templates(self) -> List[WorkflowTemplateSummary]:
        root = self.templates_root
        if not root.exists():
            return []
        summaries: List[WorkflowTemplateSummary] = []
        for child in sorted(root.iterdir()):
            if not child.is_dir():
                continue
            meta_path = child / "meta.json"
            if not meta_path.exists():
                continue
            meta = _read_json(meta_path)
            summaries.append(
                WorkflowTemplateSummary(
                    id=meta.get("id", child.name),
                    name=meta.get("name", child.name),
                    description=meta.get("description", ""),
                )
            )
        return summaries

    def get_workflow_variables(self, project_id: Optional[str] = None) -> WorkflowVariablesResponse:
        path = settings.AGENT_WORKFLOW_VARIABLES_FILE
        if path.exists():
            raw = _read_json(path)
        else:
            raw = {
                "builtIn": [
                    {"ref": "workflow.goal", "label": "工作流目标文本", "description": ""},
                    {"ref": "workflow.boundWorldbook", "label": "绑定世界书摘要", "description": ""},
                    {"ref": "workflow.boundCharacter", "label": "绑定角色卡摘要", "description": ""},
                ],
                "dynamicSuffixes": [
                    {"suffix": ".output", "labelPattern": "{displayName} · 上轮产物"},
                    {"suffix": ".entryDraft", "labelPattern": "{displayName} · 条目草稿"},
                ],
            }
        built_in = [
            WorkflowVariableDef(**item) for item in raw.get("builtIn", [])
        ]
        dynamic: List[WorkflowVariableDef] = []
        suffixes = raw.get("dynamicSuffixes") or [
            {"suffix": ".output", "labelPattern": "{displayName} · 世界书条目"},
        ]
        if project_id:
            try:
                project = self.get_project(project_id)
                for node in project.pipeline.nodes:
                    if not node.enabled:
                        continue
                    if node.skillId != "studio.worldbook_entry":
                        continue
                    for suffix_def in suffixes:
                        suffix = suffix_def.get("suffix", ".output")
                        if suffix != ".output":
                            continue
                        pattern = suffix_def.get(
                            "labelPattern", "{displayName} · 世界书条目"
                        )
                        ref = f"{node.id}{suffix}"
                        label = pattern.replace("{displayName}", node.displayName)
                        dynamic.append(
                            WorkflowVariableDef(ref=ref, label=label, description="")
                        )
            except FileNotFoundError:
                pass
        return WorkflowVariablesResponse(builtIn=built_in, dynamic=dynamic)

    def get_skill_templates(self) -> Dict[str, Any]:
        path = settings.AGENT_SKILL_TEMPLATES_FILE
        if not path.exists():
            raise FileNotFoundError("skill_templates.json not found")
        return _read_json(path)

    def get_niches(self) -> Dict[str, Any]:
        path = settings.AGENT_NICHES_FILE
        if not path.exists():
            return {"niches": []}
        return _read_json(path)

    def _unique_project_id(self, base_id: str) -> str:
        candidate = base_id
        n = 1
        while self._project_dir(candidate).exists():
            candidate = f"{base_id}-{n}"
            n += 1
        return candidate

    def create_project(self, req: CreateStudioProjectRequest) -> StudioProject:
        template_id = req.template_id or DEFAULT_TEMPLATE_ID
        template_dir = self.templates_root / template_id
        if not template_dir.exists():
            raise FileNotFoundError(f"Studio template not found: {template_id}")

        base_id = req.project_id or _slugify(req.name)
        project_id = self._unique_project_id(base_id)
        dest = self._project_dir(project_id)
        dest.mkdir(parents=True, exist_ok=False)

        template_meta = _read_json(template_dir / "meta.json")
        template_pipeline = _read_json(template_dir / "pipeline.json")
        now = datetime.now().isoformat()

        meta = {
            "id": project_id,
            "name": req.name,
            "description": template_meta.get("description", ""),
            "templateId": template_id,
            "characterId": None,
            "worldbookId": None,
            "createdAt": now,
            "updatedAt": now,
        }
        _write_json(dest / "meta.json", meta)
        _write_json(dest / "pipeline.json", _normalize_pipeline_dict(template_pipeline))
        return self.get_project(project_id)

    def delete_project(self, project_id: str) -> None:
        project_dir = self._project_dir(project_id)
        if not project_dir.exists():
            raise FileNotFoundError(f"Studio project not found: {project_id}")
        shutil.rmtree(project_dir)
        runs_dir = settings.AGENT_STUDIO_RUNS_PATH / project_id
        if runs_dir.exists():
            shutil.rmtree(runs_dir)

    def ensure_default_project(self) -> None:
        """Copy example template into default project if missing."""
        default_dir = self._project_dir("default")
        if default_dir.exists():
            return
        template_dir = self.templates_root / DEFAULT_TEMPLATE_ID
        if not template_dir.exists():
            logger.warning("builtin.studio.example template missing; skip default project seed")
            return
        default_dir.mkdir(parents=True, exist_ok=True)
        template_meta = _read_json(template_dir / "meta.json")
        now = datetime.now().isoformat()
        meta = {
            "id": "default",
            "name": "示例角色项目",
            "description": template_meta.get("description", ""),
            "templateId": DEFAULT_TEMPLATE_ID,
            "characterId": None,
            "worldbookId": None,
            "createdAt": now,
            "updatedAt": now,
        }
        _write_json(default_dir / "meta.json", meta)
        shutil.copy2(template_dir / "pipeline.json", default_dir / "pipeline.json")


studio_project_service = StudioProjectService()

try:
    studio_project_service.ensure_default_project()
except Exception as _seed_err:
    logger.warning("Studio default project seed skipped: %s", _seed_err)