SillyTavern_replica/backend/services/character_service.py

"""
角色卡服务 - 严格按照 internal.py 的数据结构
每个角色一个文件夹，包含 character.json、avatar.png 和 chats/
"""
import json
from pathlib import Path
from typing import List, Optional
from datetime import datetime
import uuid

try:
    from backend.models.internal import CharacterCard
    from backend.core.config import settings
    from backend.services.character_card_converter import CharacterCardConverter
except ImportError:
    from models.internal import CharacterCard
    from core.config import settings
    from services.character_card_converter import CharacterCardConverter


class CharacterService:
    """角色卡管理服务"""
    
    def __init__(self):
        self.characters_dir = settings.CHARACTERS_PATH
        self.converter = CharacterCardConverter()
        
        # 确保目录存在
        self.characters_dir.mkdir(parents=True, exist_ok=True)
    
    def scan_all_characters(self) -> List[CharacterCard]:
        """
        扫描所有角色卡
        
        Returns:
            按 lastChatAt 排序的角色卡列表（最新的在前）
        """
        characters = []
        
        for char_folder in self.characters_dir.iterdir():
            if not char_folder.is_dir():
                continue
            
            try:
                character = self._load_character_from_folder(char_folder)
                if character:
                    characters.append(character)
            except Exception as e:
                print(f"加载角色卡失败 {char_folder.name}: {e}")
                continue
        
        # 按最后聊天时间排序（None 排最后）
        characters.sort(
            key=lambda c: c.lastChatAt or 0, 
            reverse=True
        )
        
        return characters
    
    def _load_character_from_folder(self, folder: Path) -> Optional[CharacterCard]:
        """
        从文件夹加载角色卡
        
        Args:
            folder: 角色文件夹路径
        
        Returns:
            CharacterCard 对象或 None
        """
        # 1. 读取 character.json（必须存在）
        char_file = folder / "character.json"
        if not char_file.exists():
            return None
        
        with open(char_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # 2. 检查是否有 avatar.png
        avatar_path = None
        avatar_file = folder / "avatar.png"
        if avatar_file.exists():
            # 存储相对路径，用于前端访问
            avatar_path = f"/api/characters/{folder.name}/avatar"
        
        # 3. 计算最后聊天时间
        last_chat_at = self._get_last_chat_timestamp(folder)
        
        # 4. 构建 CharacterCard 对象（严格按照数据结构）
        character = CharacterCard(
            id=data.get('id', str(uuid.uuid4())),
            name=data['name'],
            description=data.get('description', ''),
            personality=data.get('personality', ''),
            scenario=data.get('scenario', ''),
            first_mes=data.get('first_mes', ''),
            mes_example=data.get('mes_example', ''),
            categories=data.get('categories', []),
            tags=data.get('tags', []),  # ✅ 使用标签数组
            worldInfoId=data.get('worldInfoId'),
            outputSchema=data.get('outputSchema'),
            avatarPath=avatar_path,
            alternate_greetings=data.get('alternate_greetings', []),
            tableMaintenancePrompt=data.get('tableMaintenancePrompt'),
            imageGenerationPrompt=data.get('imageGenerationPrompt'),
            tableHeaders=data.get('tableHeaders'),  # ✅ 动态表格表头
            tableDefaults=data.get('tableDefaults'),  # ✅ 动态表格默认值
            createdAt=data.get('createdAt', int(datetime.now().timestamp())),
            updatedAt=data.get('updatedAt', int(datetime.now().timestamp())),
            lastChatAt=last_chat_at,
            isFavorite=data.get('isFavorite', False),
            version=data.get('version', 1)
        )
        
        return character
    
    def _get_last_chat_timestamp(self, char_folder: Path) -> Optional[int]:
        """
        获取角色的最后聊天时间戳
        
        通过扫描 chats 目录下所有 .jsonl 文件的修改时间
        """
        chats_dir = char_folder / "chats"
        if not chats_dir.exists():
            return None
        
        latest_time = None
        
        for chat_file in chats_dir.glob("*.jsonl"):
            file_mtime = int(chat_file.stat().st_mtime)
            if latest_time is None or file_mtime > latest_time:
                latest_time = file_mtime
        
        return latest_time
    
    def get_character_by_name(self, name: str) -> Optional[CharacterCard]:
        """根据角色名获取角色卡"""
        char_folder = self.characters_dir / name
        if not char_folder.exists():
            return None
        
        return self._load_character_from_folder(char_folder)
    
    def create_character(self, character_data: dict) -> CharacterCard:
        """
        创建新角色卡
        
        Args:
            character_data: 角色数据字典
        
        Returns:
            创建的 CharacterCard 对象
        """
        # 生成唯一ID
        if 'id' not in character_data:
            character_data['id'] = str(uuid.uuid4())
        
        # 设置时间戳
        now = int(datetime.now().timestamp())
        character_data['createdAt'] = now
        character_data['updatedAt'] = now
        character_data['lastChatAt'] = None
        
        # 创建文件夹
        char_name = character_data['name']
        char_folder = self.characters_dir / char_name
        char_folder.mkdir(parents=True, exist_ok=True)
        
        # 创建 chats 目录
        chats_dir = char_folder / "chats"
        chats_dir.mkdir(exist_ok=True)
        
        # 保存 character.json
        char_file = char_folder / "character.json"
        with open(char_file, 'w', encoding='utf-8') as f:
            json.dump(character_data, f, ensure_ascii=False, indent=2)
        
        return self._load_character_from_folder(char_folder)
    
    def update_character(self, name: str, updates: dict) -> CharacterCard:
        """
        更新角色卡
        
        Args:
            name: 角色名（旧名称，用于定位文件夹）
            updates: 更新的字段（可以包含 name 字段来重命名）
        
        Returns:
            更新后的 CharacterCard 对象
        """
        char_folder = self.characters_dir / name
        char_file = char_folder / "character.json"
        
        if not char_file.exists():
            raise FileNotFoundError(f"角色卡不存在: {name}")
        
        # 读取现有数据
        with open(char_file, 'r', encoding='utf-8') as f:
            existing_data = json.load(f)
        
        # 检查是否需要重命名
        new_name = updates.get('name')
        needs_rename = new_name and new_name != name
        
        if needs_rename:
            # 验证新名称是否合法
            if not new_name or new_name.strip() == '':
                raise ValueError("角色名不能为空")
            
            # 检查新名称是否已存在
            new_folder = self.characters_dir / new_name
            if new_folder.exists():
                raise FileExistsError(f"角色 '{new_name}' 已存在")
            
            # 重命名文件夹
            try:
                import shutil
                shutil.move(str(char_folder), str(new_folder))
                char_folder = new_folder
                char_file = char_folder / "character.json"
            except Exception as e:
                raise RuntimeError(f"重命名文件夹失败: {str(e)}")
        
        # 合并更新
        existing_data.update(updates)
        existing_data['updatedAt'] = int(datetime.now().timestamp())
        
        # 保存
        with open(char_file, 'w', encoding='utf-8') as f:
            json.dump(existing_data, f, ensure_ascii=False, indent=2)
        
        return self._load_character_from_folder(char_folder)
    
    def delete_character(self, name: str) -> bool:
        """
        删除角色卡（包括所有聊天记录）
        
        Args:
            name: 角色名
        
        Returns:
            是否成功删除
        """
        char_folder = self.characters_dir / name
        if not char_folder.exists():
            return False
        
        import shutil
        shutil.rmtree(char_folder)
        return True
    
    def save_avatar(self, name: str, image_data: bytes) -> str:
        """
        保存角色头像
        
        Args:
            name: 角色名
            image_data: 图片二进制数据
        
        Returns:
            头像访问路径
        """
        char_folder = self.characters_dir / name
        avatar_file = char_folder / "avatar.png"
        
        with open(avatar_file, 'wb') as f:
            f.write(image_data)
        
        return f"/api/characters/{name}/avatar"
    
    def import_from_png(self, png_data: bytes, filename: str) -> CharacterCard:
        """
        从 SillyTavern PNG 导入角色卡
        
        Args:
            png_data: PNG 文件二进制数据
            filename: 原始文件名
        
        Returns:
            创建的 CharacterCard 对象
        """
        # 1. 提取嵌入数据
        st_data = self.converter.extract_from_png(png_data)
        if not st_data:
            raise ValueError("PNG文件中没有嵌入角色数据")
        
        # 2. 转换为内部格式
        character = self.converter.st_to_internal(st_data)
        
        # 3. 创建角色文件夹
        char_name = character.name
        char_folder = self.characters_dir / char_name
        char_folder.mkdir(parents=True, exist_ok=True)
        
        # 4. 保存 PNG 作为 avatar.png
        avatar_file = char_folder / "avatar.png"
        with open(avatar_file, 'wb') as f:
            f.write(png_data)
        
        # 5. 保存 character.json
        char_file = char_folder / "character.json"
        with open(char_file, 'w', encoding='utf-8') as f:
            json.dump(character.dict(), f, ensure_ascii=False, indent=2)
        
        # 6. 创建 chats 目录
        (char_folder / "chats").mkdir(exist_ok=True)
        
        return character
    
    def export_as_png(self, name: str) -> bytes:
        """
        导出角色为 SillyTavern PNG 格式
        
        Args:
            name: 角色名
        
        Returns:
            PNG 文件二进制数据
        """
        character = self.get_character_by_name(name)
        if not character:
            raise FileNotFoundError(f"角色 '{name}' 不存在")
        
        # 获取头像路径
        avatar_path = None
        if character.avatarPath:
            # 从路径中提取文件名
            avatar_filename = character.avatarPath.split('/')[-1].split('?')[0]
            char_folder = self.characters_dir / name
            avatar_file = char_folder / avatar_filename
            if avatar_file.exists():
                avatar_path = str(avatar_file)
        
        # 如果没有头像，使用默认图片
        use_default = False
        if not avatar_path:
            default_avatar = self.characters_dir / "defult.png"
            if default_avatar.exists():
                avatar_path = str(default_avatar)
                use_default = True
                print(f"使用默认头像: {avatar_path}")
            else:
                print("警告: 没有找到默认头像")
        
        # 生成 PNG
        return self.converter.export_as_png(character, avatar_path, use_default_avatar=use_default)