kotones-auto-assistant/kotonebot/kaa/config/produce.py

import os
import json
import uuid
import re
import logging
from typing import Literal
from pydantic import BaseModel, ConfigDict

from .const import ProduceAction, RecommendCardDetectionMode

logger = logging.getLogger(__name__)

class ConfigBaseModel(BaseModel):
    model_config = ConfigDict(use_attribute_docstrings=True)


class ProduceData(ConfigBaseModel):
    enabled: bool = False
    """是否启用培育"""
    mode: Literal['regular', 'pro', 'master'] = 'regular'
    """
    培育模式。
    进行一次 REGULAR 培育需要 ~30min，进行一次 PRO 培育需要 ~1h（具体视设备性能而定）。
    """
    produce_count: int = 1
    """培育的次数。"""
    idols: list[str] = []
    """
    要培育偶像的 IdolCardSkin.id。将会按顺序循环选择培育。
    """
    memory_sets: list[int] = []
    """要使用的回忆编成编号，从 1 开始。将会按顺序循环选择使用。"""
    support_card_sets: list[int] = []
    """要使用的支援卡编成编号，从 1 开始。将会按顺序循环选择使用。"""
    auto_set_memory: bool = False
    """是否自动编成回忆。此选项优先级高于回忆编成编号。"""
    auto_set_support_card: bool = False
    """是否自动编成支援卡。此选项优先级高于支援卡编成编号。"""
    use_pt_boost: bool = False
    """是否使用支援强化 Pt 提升。"""
    use_note_boost: bool = False
    """是否使用笔记数提升。"""
    follow_producer: bool = False
    """是否关注租借了支援卡的制作人。"""
    self_study_lesson: Literal['dance', 'visual', 'vocal'] = 'dance'
    """自习课类型。"""
    prefer_lesson_ap: bool = False
    """
    优先 SP 课程。

    启用后，若出现 SP 课程，则会优先执行 SP 课程，而不是推荐课程。
    若出现多个 SP 课程，随机选择一个。
    """
    actions_order: list[ProduceAction] = [
        ProduceAction.RECOMMENDED,
        ProduceAction.VISUAL,
        ProduceAction.VOCAL,
        ProduceAction.DANCE,
        ProduceAction.ALLOWANCE,
        ProduceAction.OUTING,
        ProduceAction.STUDY,
        ProduceAction.CONSULT,
        ProduceAction.REST,
    ]
    """
    行动优先级

    每一周的行动将会按这里设置的优先级执行。
    """
    recommend_card_detection_mode: RecommendCardDetectionMode = RecommendCardDetectionMode.NORMAL
    """
    推荐卡检测模式

    严格模式下，识别速度会降低，但识别准确率会提高。
    """
    use_ap_drink: bool = False
    """
    AP 不足时自动使用 AP 饮料
    """
    skip_commu: bool = True
    """检测并跳过交流"""


class ProduceSolution(ConfigBaseModel):
    """培育方案"""
    id: str
    """方案唯一标识符"""
    name: str
    """方案名称"""
    description: str | None = None
    """方案描述"""
    data: ProduceData
    """培育数据"""


class ProduceSolutionManager:
    """培育方案管理器"""

    SOLUTIONS_DIR = "conf/produce"

    def __init__(self):
        """初始化管理器，确保目录存在"""
        os.makedirs(self.SOLUTIONS_DIR, exist_ok=True)

    def _sanitize_filename(self, name: str) -> str:
        """
        清理文件名中的非法字符

        :param name: 原始名称
        :return: 清理后的文件名
        """
        # 替换 \/:*?"<>| 为下划线
        return re.sub(r'[\\/:*?"<>|]', '_', name)

    def _get_file_path(self, name: str) -> str:
        """
        根据方案名称获取文件路径

        :param name: 方案名称
        :return: 文件路径
        """
        safe_name = self._sanitize_filename(name)
        return os.path.join(self.SOLUTIONS_DIR, f"{safe_name}.json")

    def new(self, name: str) -> ProduceSolution:
        """
        创建新的培育方案

        :param name: 方案名称
        :return: 新创建的方案
        """
        solution = ProduceSolution(
            id=uuid.uuid4().hex,
            name=name,
            data=ProduceData()
        )
        return solution

    def list(self) -> list[ProduceSolution]:
        """
        列出所有培育方案

        :return: 方案列表
        """
        solutions = []
        if not os.path.exists(self.SOLUTIONS_DIR):
            return solutions

        for filename in os.listdir(self.SOLUTIONS_DIR):
            if filename.endswith('.json'):
                try:
                    file_path = os.path.join(self.SOLUTIONS_DIR, filename)
                    with open(file_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                    solution = ProduceSolution.model_validate_json(data)
                    solutions.append(solution)
                    logger.info(f"Loaded produce solution from {file_path}")
                except Exception:
                    logger.warning(f"Failed to load produce solution from {file_path}")
                    continue

        return solutions

    def delete(self, id: str) -> None:
        """
        删除指定ID的培育方案

        :param id: 方案ID
        """
        if not os.path.exists(self.SOLUTIONS_DIR):
            return

        for filename in os.listdir(self.SOLUTIONS_DIR):
            if filename.endswith('.json'):
                try:
                    file_path = os.path.join(self.SOLUTIONS_DIR, filename)
                    with open(file_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                    if data.get('id') == id:
                        os.remove(file_path)
                        return
                except Exception:
                    continue

    def save(self, id: str, solution: ProduceSolution) -> None:
        """
        保存培育方案

        :param id: 方案ID
        :param solution: 方案对象
        """
        # 确保ID一致
        solution.id = id

        file_path = self._get_file_path(solution.name)
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(solution.model_dump_json(indent=4), f, ensure_ascii=False)

    def read(self, id: str) -> ProduceSolution:
        """
        读取指定ID的培育方案

        :param id: 方案ID
        :return: 方案对象
        :raises FileNotFoundError: 当方案不存在时
        """
        if not os.path.exists(self.SOLUTIONS_DIR):
            raise FileNotFoundError(f"Solution with id '{id}' not found")

        for filename in os.listdir(self.SOLUTIONS_DIR):
            if filename.endswith('.json'):
                try:
                    file_path = os.path.join(self.SOLUTIONS_DIR, filename)
                    with open(file_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                    if data.get('id') == id:
                        return ProduceSolution.model_validate_json(data)
                except Exception:
                    continue

        raise FileNotFoundError(f"Solution with id '{id}' not found")