cmshoppe/app/prompts.py

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

from __future__ import annotations

import os
from importlib import resources

from . import appconfig

TITLE_PROMPT_PATH = appconfig.title_prompt_path()
TITLE_TEMPLATES_DIR = appconfig.title_templates_dir()
COVER_PROMPTS_DIR = appconfig.cover_prompts_dir()
IMAGE_STUDIO_PROMPTS_DIR = appconfig.image_studio_prompts_dir()
TEMPLATE_EXT = ".txt"
INVALID_NAME_CHARS = set('\\/:*?"<>|')
DEFAULT_PROMPTS_PACKAGE = "app.default_prompts"
DEFAULT_COVER_PROMPTS_PACKAGE = "app.default_prompts.cover"


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 ensure_default_prompts(
    title_prompt_path=TITLE_PROMPT_PATH,
    cover_prompts_dir=COVER_PROMPTS_DIR,
    title_templates_dir=None,
) -> None:
    """Seed bundled default prompts into an empty user data directory.

    Existing user prompts are never overwritten. Defaults live in the packaged
    app and are copied into data/ only when the corresponding user files are
    missing, or the title prompt file is blank.
    """

    if not _has_non_empty_file(title_prompt_path):
        default_title = _read_default_prompt_text("title_prompt.txt")
        if default_title:
            save_title_prompt(default_title, title_prompt_path)

    if title_templates_dir is None:
        title_templates_dir = _title_templates_dir_for_prompt_path(title_prompt_path)
    if not list_templates(title_templates_dir):
        default_title = _read_default_prompt_text("title_prompt.txt")
        if default_title:
            save_template("默认", default_title, title_templates_dir)

    if not list_cover_templates(cover_prompts_dir):
        for name, text in _iter_default_cover_templates():
            save_cover_template(name, text, cover_prompts_dir)


def list_templates(directory):
    """Return prompt 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_template(name, directory) -> str:
    """Load one 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_template(name, text, directory) -> None:
    """Save one 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_template(old, new, directory) -> None:
    """Rename a 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_template(name, directory) -> None:
    """Delete one prompt template."""

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


def list_title_templates(directory=TITLE_TEMPLATES_DIR):
    """Return title prompt template names sorted by display name."""

    return list_templates(directory)


def load_title_template(name, directory=TITLE_TEMPLATES_DIR) -> str:
    """Load one title prompt template."""

    return load_template(name, directory)


def save_title_template(name, text, directory=TITLE_TEMPLATES_DIR) -> None:
    """Save one title prompt template as UTF-8."""

    save_template(name, text, directory)


def rename_title_template(old, new, directory=TITLE_TEMPLATES_DIR) -> None:
    """Rename a title prompt template with duplicate-name protection."""

    rename_template(old, new, directory)


def delete_title_template(name, directory=TITLE_TEMPLATES_DIR) -> None:
    """Delete one title prompt template."""

    delete_template(name, directory)


def list_image_studio_templates(directory=IMAGE_STUDIO_PROMPTS_DIR):
    """Return AI studio prompt template names sorted by display name."""

    return list_templates(directory)


def load_image_studio_template(name, directory=IMAGE_STUDIO_PROMPTS_DIR) -> str:
    """Load one AI studio prompt template as raw text."""

    return load_template(name, directory)


def save_image_studio_template(name, text, directory=IMAGE_STUDIO_PROMPTS_DIR) -> None:
    """Save one AI studio prompt template as raw UTF-8 text."""

    save_template(name, text, directory)


def rename_image_studio_template(old, new, directory=IMAGE_STUDIO_PROMPTS_DIR) -> None:
    """Rename an AI studio prompt template with duplicate-name protection."""

    rename_template(old, new, directory)


def delete_image_studio_template(name, directory=IMAGE_STUDIO_PROMPTS_DIR) -> None:
    """Delete one AI studio prompt template."""

    delete_template(name, directory)


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

    return list_templates(directory)


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

    return load_template(name, directory)


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

    save_template(name, text, directory)


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

    rename_template(old, new, directory)


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

    delete_template(name, directory)


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


def _title_templates_dir_for_prompt_path(title_prompt_path) -> str:
    data_root = os.path.dirname(os.path.abspath(title_prompt_path))
    return os.path.join(data_root, "prompts", "title")


def _has_non_empty_file(path) -> bool:
    try:
        if not os.path.exists(path):
            return False
        with open(path, "r", encoding="utf-8") as fh:
            return bool(fh.read().strip())
    except OSError:
        return False


def _read_default_prompt_text(filename) -> str:
    try:
        return (
            resources.files(DEFAULT_PROMPTS_PACKAGE)
            .joinpath(filename)
            .read_text(encoding="utf-8")
        )
    except (FileNotFoundError, ModuleNotFoundError, OSError):
        return ""


def _iter_default_cover_templates():
    try:
        root = resources.files(DEFAULT_COVER_PROMPTS_PACKAGE)
    except (ModuleNotFoundError, OSError):
        return
    for entry in root.iterdir():
        if not entry.is_file() or not entry.name.lower().endswith(TEMPLATE_EXT):
            continue
        name = entry.name[: -len(TEMPLATE_EXT)]
        try:
            text = entry.read_text(encoding="utf-8")
        except OSError:
            continue
        yield name, text