cmshoppe/app/prompts.py

"""Prompt file management and variable rendering."""

from __future__ import annotations

import os

from . import appconfig

TITLE_PROMPT_PATH = appconfig.title_prompt_path()
COVER_PROMPTS_DIR = appconfig.cover_prompts_dir()
TEMPLATE_EXT = ".txt"
INVALID_NAME_CHARS = set('\\/:*?"<>|')


class PromptError(RuntimeError):
    """Raised when prompt files or template names are invalid."""


def load_title_prompt(path=TITLE_PROMPT_PATH) -> str:
    """Load the title prompt text. Missing file means an empty prompt."""

    if not os.path.exists(path):
        return ""
    with open(path, "r", encoding="utf-8") as fh:
        return fh.read()


def save_title_prompt(text, path=TITLE_PROMPT_PATH) -> None:
    """Save the title prompt text as UTF-8."""

    directory = os.path.dirname(os.path.abspath(path))
    if directory:
        os.makedirs(directory, exist_ok=True)
    with open(path, "w", encoding="utf-8") as fh:
        fh.write(str(text or ""))


def list_cover_templates(directory=COVER_PROMPTS_DIR):
    """Return cover template names sorted by display name."""

    if not os.path.isdir(directory):
        return []
    names = []
    for filename in os.listdir(directory):
        if filename.lower().endswith(TEMPLATE_EXT):
            names.append(filename[: -len(TEMPLATE_EXT)])
    return sorted(names, key=str.casefold)


def load_cover_template(name, directory=COVER_PROMPTS_DIR) -> str:
    """Load one cover prompt template."""

    path = _template_path(name, directory)
    if not os.path.exists(path):
        raise PromptError(f"封面提示词模板不存在: {_normalize_name(name)}")
    with open(path, "r", encoding="utf-8") as fh:
        return fh.read()


def save_cover_template(name, text, directory=COVER_PROMPTS_DIR) -> None:
    """Save one cover prompt template as UTF-8."""

    path = _template_path(name, directory)
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", encoding="utf-8") as fh:
        fh.write(str(text or ""))


def rename_cover_template(old, new, directory=COVER_PROMPTS_DIR) -> None:
    """Rename a cover prompt template with duplicate-name protection."""

    old_path = _template_path(old, directory)
    new_path = _template_path(new, directory)
    if not os.path.exists(old_path):
        raise PromptError(f"封面提示词模板不存在: {_normalize_name(old)}")
    if os.path.exists(new_path):
        raise PromptError(f"封面提示词模板已存在: {_normalize_name(new)}")
    os.makedirs(os.path.dirname(new_path), exist_ok=True)
    os.replace(old_path, new_path)


def delete_cover_template(name, directory=COVER_PROMPTS_DIR) -> None:
    """Delete one cover prompt template."""

    path = _template_path(name, directory)
    if not os.path.exists(path):
        raise PromptError(f"封面提示词模板不存在: {_normalize_name(name)}")
    os.remove(path)


def render_prompt(template_text, task) -> str:
    """Render known task variables in a prompt template."""

    values = {
        "旧标题": _task_value(task, "old_title", "旧标题"),
        "新标题": _task_value(task, "new_title", "新标题"),
        "商品id": _task_value(task, "item_id", "商品id", "商品ID"),
        "店铺": _task_value(task, "shop", "account_name", "店铺", "alias"),
    }
    rendered = str(template_text or "")
    for name, value in values.items():
        rendered = rendered.replace("{" + name + "}", value)
    return rendered


def _task_value(task, *names) -> str:
    for name in names:
        value = None
        if isinstance(task, dict):
            value = task.get(name)
        else:
            value = getattr(task, name, None)
        if value is not None:
            return str(value)
    return ""


def _template_path(name, directory) -> str:
    normalized = _normalize_name(name)
    return os.path.abspath(os.path.join(directory, normalized + TEMPLATE_EXT))


def _normalize_name(name) -> str:
    value = str(name or "").strip()
    if value.lower().endswith(TEMPLATE_EXT):
        value = value[: -len(TEMPLATE_EXT)]
    value = value.strip()
    if not value:
        raise PromptError("封面提示词模板名不能为空")
    if value in {".", ".."} or any(char in INVALID_NAME_CHARS for char in value):
        raise PromptError(f"封面提示词模板名非法: {value}")
    if os.path.basename(value) != value:
        raise PromptError(f"封面提示词模板名非法: {value}")
    return value