SillyTavern_replica/backend/services/chat_service.py

"""
聊天服务 - 处理聊天记录的读写操作

基于 SillyTavern JSONL 格式的聊天记录管理
"""
from pathlib import Path
from typing import Dict, List, Optional, Any
import json
import logging
from datetime import datetime
import uuid

from core.config import settings

logger = logging.getLogger(__name__)


class ChatService:
    """聊天服务类，处理聊天记录的CRUD操作"""
    
    def __init__(self, data_path: Path):
        """
        初始化聊天服务
        
        Args:
            data_path: 数据目录路径
        """
        self.data_path = data_path
        self.chat_dir = data_path / "chat"
        self.chat_dir.mkdir(parents=True, exist_ok=True)
    
    def list_all_chats(self) -> Dict[str, List[Dict]]:
        """
        获取所有角色和聊天列表
        
        Returns:
            Dict[str, List[Dict]]: 字典结构，键是角色名称，值是该角色的聊天信息列表
        """
        result = {}
        
        if not self.chat_dir.exists():
            logger.warning(f"聊天目录不存在: {self.chat_dir}")
            return result
        
        for role_dir in self.chat_dir.iterdir():
            try:
                if role_dir.is_dir():
                    chats = []
                    
                    for chat_file in role_dir.glob("*.jsonl"):
                        chat_info = self._get_chat_summary(role_dir.name, chat_file.stem)
                        if chat_info:
                            chats.append(chat_info)
                    
                    if chats:
                        result[role_dir.name] = chats
                        
            except Exception as e:
                logger.error(f"处理角色目录 {role_dir.name} 时出错: {str(e)}")
                continue
        
        return result
    
    def _get_chat_summary(self, role_name: str, chat_name: str) -> Optional[Dict]:
        """
        获取聊天摘要信息
        
        Args:
            role_name: 角色名称
            chat_name: 聊天名称
            
        Returns:
            Dict: 聊天摘要信息，如果文件不存在则返回None
        """
        chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
        
        if not chat_file.exists():
            return None
        
        try:
            with open(chat_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            if not lines:
                return None
            
            # 第一行是header
            header = json.loads(lines[0])
            
            # 计算消息数量（排除header）
            message_count = len(lines) - 1
            
            # 获取最后修改时间
            last_modified = datetime.fromtimestamp(
                chat_file.stat().st_mtime
            ).isoformat()
            
            # 获取最后一条消息预览
            last_message = ""
            if message_count > 0:
                try:
                    last_msg_data = json.loads(lines[-1])
                    last_message = last_msg_data.get("mes", "")
                except:
                    pass
            
            return {
                "chat_name": chat_name,
                "user_name": header.get("user_name", "User"),
                "character_name": header.get("character_name", ""),
                "last_modified": last_modified,
                "message_count": message_count,
                "last_message": last_message
            }
            
        except Exception as e:
            logger.error(f"读取聊天摘要失败 {role_name}/{chat_name}: {str(e)}")
            return None
    
    def get_chat(self, role_name: str, chat_name: str) -> Dict[str, Any]:
        """
        获取指定聊天的完整内容
        
        Args:
            role_name: 角色名称
            chat_name: 聊天名称
            
        Returns:
            Dict: 包含metadata和messages的字典
            
        Raises:
            FileNotFoundError: 聊天文件不存在
        """
        chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
        
        if not chat_file.exists():
            raise FileNotFoundError(f"Chat not found: {role_name}/{chat_name}")
        
        try:
            with open(chat_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            if not lines:
                raise ValueError(f"Empty chat file: {role_name}/{chat_name}")
            
            # 第一行是header
            header = json.loads(lines[0])
            
            # 解析消息
            messages = []
            for i, line in enumerate(lines[1:], start=1):
                if line.strip():  # 跳过空行
                    msg_data = json.loads(line)
                    # 确保有floor字段
                    if "floor" not in msg_data:
                        msg_data["floor"] = i
                    
                    messages.append(msg_data)
            
            return {
                "header": header,  # 完整的 header，包含 tableHeaders, tableDefaults, tableData
                "metadata": {
                    "user_name": header.get("user_name", "User"),
                    "character_name": header.get("character_name", ""),
                    "chat_id": header.get("chat_id_hash", ""),
                    "integrity": header.get("integrity", "")
                },
                "messages": messages
            }
            
        except Exception as e:
            logger.error(f"读取聊天失败 {role_name}/{chat_name}: {str(e)}")
            raise
    
    def get_message(self, role_name: str, chat_name: str, floor: int) -> Dict:
        """
        获取指定楼层的消息
        
        Args:
            role_name: 角色名称
            chat_name: 聊天名称
            floor: 楼层号
            
        Returns:
            Dict: 消息数据，如果不存在则返回 None
        """
        chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
        
        if not chat_file.exists():
            return None
        
        try:
            with open(chat_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            # 找到对应的消息行（floor + 1，因为第0行是header）
            message_line_index = floor + 1
            
            if message_line_index >= len(lines):
                return None
            
            # 解析并返回消息
            msg_data = json.loads(lines[message_line_index])
            return msg_data
            
        except Exception as e:
            logger.error(f"获取消息失败 {role_name}/{chat_name}/{floor}: {str(e)}")
            return None
    
    def create_chat(self, role_name: str, chat_name: str, metadata: Dict = None) -> Dict:
        """
        创建新聊天
        
        Args:
            role_name: 角色名称
            chat_name: 聊天名称
            metadata: 聊天元数据
            
        Returns:
            Dict: 创建的聊天信息
            
        Raises:
            FileExistsError: 聊天已存在
        """
        chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
        
        if chat_file.exists():
            raise FileExistsError(f"Chat already exists: {role_name}/{chat_name}")
        
        # 创建角色目录
        chat_file.parent.mkdir(parents=True, exist_ok=True)
        
        # 尝试从角色卡获取 tags（关键字列表）
        tags = []
        
        try:
            character_file = settings.CHARACTERS_PATH / role_name / "character.json"
            if character_file.exists():
                with open(character_file, 'r', encoding='utf-8') as f:
                    character_data = json.load(f)
                
                tags = character_data.get('tags', [])
                
                logger.info(f"从角色卡 {role_name} 继承标签: {tags}")
        except Exception as e:
            logger.warning(f"读取角色卡失败，使用空标签: {e}")
        
        # 构建header
        header = {
            "user_name": metadata.get("user_name", "User") if metadata else "User",
            "character_name": metadata.get("character_name", role_name) if metadata else role_name,
            "integrity": str(uuid.uuid4()),
            "chat_id_hash": str(uuid.uuid4()),
            "note_prompt": "",
            "note_interval": 0,
            "note_position": 0,
            "note_depth": 0,
            "note_role": 0,
            "extensions": {},
            "timedWorldInfo": {},
            "variables": {},
            "tainted": False,
            "lastInContextMessageId": -1,
            "tags": tags  # ✅ 使用标签数组替代 tableHeaders/tableDefaults/tableData
        }
        
        # 写入header
        with open(chat_file, 'w', encoding='utf-8') as f:
            f.write(json.dumps(header, ensure_ascii=False) + '\n')
        
        return {
            "role_name": role_name,
            "chat_name": chat_name,
            "metadata": {
                "user_name": header["user_name"],
                "character_name": header["character_name"]
            }
        }
    
    def add_message(self, role_name: str, chat_name: str, message_data: Dict) -> Dict:
        """
        向聊天添加新消息
        
        Args:
            role_name: 角色名称
            chat_name: 聊天名称
            message_data: 消息数据
            
        Returns:
            Dict: 添加的消息
            
        Raises:
            FileNotFoundError: 聊天不存在
        """
        chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
        
        if not chat_file.exists():
            raise FileNotFoundError(f"Chat not found: {role_name}/{chat_name}")
        
        try:
            # 读取现有消息以确定floor
            with open(chat_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            # 计算下一个floor号
            next_floor = len(lines) - 1  # 减去header行
            
            # 构建完整的消息数据
            full_message = {
                "name": message_data.get("name", "User"),
                "is_user": message_data.get("is_user", True),
                "is_system": message_data.get("is_system", False),
                "floor": next_floor,
                "send_date": message_data.get("send_date", str(int(datetime.now().timestamp() * 1000))),
                "mes": message_data.get("mes", ""),
                "extra": message_data.get("extra", {}),
                "swipes": message_data.get("swipes", []),
                "swipe_id": message_data.get("swipe_id", 0),
                "force_avatar": None,
                "variables": [],
                "variables_initialized": [],
                "is_ejs_processed": []
            }
            
            # 追加消息到文件
            with open(chat_file, 'a', encoding='utf-8') as f:
                f.write(json.dumps(full_message, ensure_ascii=False) + '\n')
            
            return full_message
            
        except Exception as e:
            logger.error(f"添加消息失败 {role_name}/{chat_name}: {str(e)}")
            raise
    
    def update_message(self, role_name: str, chat_name: str, floor: int, update_data: Dict) -> Dict:
        """
        更新指定楼层的消息
        
        Args:
            role_name: 角色名称
            chat_name: 聊天名称
            floor: 楼层号
            update_data: 更新的数据
            
        Returns:
            Dict: 更新后的消息
            
        Raises:
            FileNotFoundError: 聊天不存在
            ValueError: 楼层不存在
        """
        chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
        
        if not chat_file.exists():
            raise FileNotFoundError(f"Chat not found: {role_name}/{chat_name}")
        
        try:
            # 读取所有行
            with open(chat_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            # 找到对应的消息行（floor + 1，因为第0行是header）
            message_line_index = floor + 1
            
            if message_line_index >= len(lines):
                raise ValueError(f"Floor {floor} not found in chat")
            
            # 解析并更新消息
            msg_data = json.loads(lines[message_line_index])
            msg_data.update(update_data)
            
            # 写回文件
            lines[message_line_index] = json.dumps(msg_data, ensure_ascii=False) + '\n'
            
            with open(chat_file, 'w', encoding='utf-8') as f:
                f.writelines(lines)
            
            return msg_data
            
        except Exception as e:
            logger.error(f"更新消息失败 {role_name}/{chat_name}/{floor}: {str(e)}")
            raise
    
    def delete_message(self, role_name: str, chat_name: str, floor: int) -> Dict:
        """
        删除指定楼层的消息
        
        Args:
            role_name: 角色名称
            chat_name: 聊天名称
            floor: 楼层号
            
        Returns:
            Dict: 被删除的消息
            
        Raises:
            FileNotFoundError: 聊天不存在
            ValueError: 楼层不存在
        """
        chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
        
        if not chat_file.exists():
            raise FileNotFoundError(f"Chat not found: {role_name}/{chat_name}")
        
        try:
            # 读取所有行
            with open(chat_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            # 找到对应的消息行
            message_line_index = floor + 1
            
            if message_line_index >= len(lines):
                raise ValueError(f"Floor {floor} not found in chat")
            
            # 保存被删除的消息
            deleted_msg = json.loads(lines[message_line_index])
            
            # 删除该行
            del lines[message_line_index]
            
            # 重新编号后续消息的floor
            for i in range(message_line_index, len(lines)):
                if lines[i].strip():  # 跳过空行
                    msg_data = json.loads(lines[i])
                    msg_data["floor"] = i - 1  # 重新计算floor
                    lines[i] = json.dumps(msg_data, ensure_ascii=False) + '\n'
            
            # 写回文件
            with open(chat_file, 'w', encoding='utf-8') as f:
                f.writelines(lines)
            
            return deleted_msg
            
        except Exception as e:
            logger.error(f"删除消息失败 {role_name}/{chat_name}/{floor}: {str(e)}")
            raise
    
    def update_table_data(self, role_name: str, chat_name: str, table_update: Dict) -> Dict:
        """
        更新标签数据（SillyTavern 关键字机制）
        
        Args:
            role_name: 角色名称
            chat_name: 聊天名称
            table_update: 包含 tags 数组的字典
            
        Returns:
            Dict: 更新后的标签数据
            
        Raises:
            FileNotFoundError: 聊天文件不存在
        """
        chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
        
        if not chat_file.exists():
            raise FileNotFoundError(f"Chat not found: {role_name}/{chat_name}")
        
        try:
            with open(chat_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            if not lines:
                raise ValueError(f"Empty chat file: {role_name}/{chat_name}")
            
            # 读取 header
            header = json.loads(lines[0])
            
            # 获取新的标签数组
            new_tags = table_update.get('tags', [])
            
            # 更新 header 中的 tags
            header['tags'] = new_tags
            
            # 写回文件
            lines[0] = json.dumps(header, ensure_ascii=False) + '\n'
            
            with open(chat_file, 'w', encoding='utf-8') as f:
                f.writelines(lines)
            
            logger.info(f"标签数据已更新: {role_name}/{chat_name}, 标签数: {len(new_tags)}")
            
            return {
                "success": True,
                "tags": new_tags,
                "tagCount": len(new_tags)
            }
            
        except Exception as e:
            logger.error(f"更新标签数据失败 {role_name}/{chat_name}: {str(e)}")
            raise
    
    def summarize_chat_messages(
        self,
        role_name: str,
        chat_name: str,
        start_floor: int,
        end_floor: int,
        summary_text: str
    ) -> bool:
        """
        总结聊天消息：清空原文，将总结放到最后一个楼层
        
        Args:
            role_name: 角色名称
            chat_name: 聊天名称
            start_floor: 总结起始楼层（1-based）
            end_floor: 总结结束楼层（1-based）
            summary_text: 总结文本
            
        Returns:
            bool: 是否成功
        """
        chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
        
        if not chat_file.exists():
            raise FileNotFoundError(f"Chat not found: {role_name}/{chat_name}")
        
        try:
            with open(chat_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            if not lines:
                raise ValueError(f"Empty chat file: {role_name}/{chat_name}")
            
            # 转换为0-based索引
            start_idx = start_floor  # header在第0行，所以L1在第1行
            end_idx = end_floor
            
            # 验证范围
            if start_idx < 1 or end_idx >= len(lines) or start_idx > end_idx:
                raise ValueError(f"Invalid floor range: {start_floor}-{end_floor}")
            
            # 处理消息
            for i in range(start_idx, end_idx + 1):
                msg_data = json.loads(lines[i])
                
                if i < end_idx:
                    # 中间楼层：清空内容
                    msg_data['mes'] = ""
                    msg_data['is_summarized'] = True
                else:
                    # 最后一个楼层：放入总结文本
                    msg_data['mes'] = summary_text
                    msg_data['is_summary'] = True
                    msg_data['summary_range'] = f"L{start_floor}-L{end_floor}"
                
                lines[i] = json.dumps(msg_data, ensure_ascii=False) + '\n'
            
            # 写回文件
            with open(chat_file, 'w', encoding='utf-8') as f:
                f.writelines(lines)
            
            logger.info(
                f"[ChatService] 总结完成: {role_name}/{chat_name}, "
                f"楼层 {start_floor}-{end_floor}"
            )
            
            return True
            
        except Exception as e:
            logger.error(f"总结聊天消息失败 {role_name}/{chat_name}: {str(e)}")
            raise
    
    def create_branch(
        self,
        role_name: str,
        chat_name: str,
        target_floor: int,
        new_chat_name: Optional[str] = None
    ) -> Dict:
        """
        创建聊天分支
        
        复制目标楼层及之前的所有内容到一个新的聊天记录
        
        Args:
            role_name: 角色名称
            chat_name: 原聊天名称
            target_floor: 目标楼层（包含该楼层及之前的内容）
            new_chat_name: 新聊天名称（可选，默认自动生成）
            
        Returns:
            Dict: {
                "success": bool,
                "new_chat_name": str,
                "message_count": int
            }
            
        Raises:
            FileNotFoundError: 聊天不存在
            ValueError: 楼层不存在
        """
        chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
        
        if not chat_file.exists():
            raise FileNotFoundError(f"Chat not found: {role_name}/{chat_name}")
        
        try:
            # 读取所有行
            with open(chat_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            if not lines:
                raise ValueError(f"Empty chat file: {role_name}/{chat_name}")
            
            # 验证目标楼层
            # floor + 1 是因为第0行是header
            message_line_index = target_floor + 1
            if message_line_index >= len(lines):
                raise ValueError(f"Floor {target_floor} not found in chat (total messages: {len(lines) - 1})")
            
            # 生成新聊天名称
            if not new_chat_name:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                new_chat_name = f"branch_{chat_name}_{timestamp}"
            
            # 创建新聊天文件
            new_chat_file = self.chat_dir / role_name / f"{new_chat_name}.jsonl"
            
            if new_chat_file.exists():
                raise FileExistsError(f"Branch chat already exists: {new_chat_name}")
            
            # 复制 header 和目标楼层及之前的消息
            branch_lines = [lines[0]]  # header
            for i in range(1, message_line_index + 1):
                msg_data = json.loads(lines[i])
                # 重新分配 floor（从0开始）
                msg_data["floor"] = i - 1
                branch_lines.append(json.dumps(msg_data, ensure_ascii=False) + '\n')
            
            # 写入新文件
            with open(new_chat_file, 'w', encoding='utf-8') as f:
                f.writelines(branch_lines)
            
            message_count = len(branch_lines) - 1  # 减去header
            
            logger.info(
                f"[ChatService] 创建分支成功: {role_name}/{chat_name} -> {new_chat_name}, "
                f"楼层: 0-{target_floor}, 消息数: {message_count}"
            )
            
            return {
                "success": True,
                "new_chat_name": new_chat_name,
                "message_count": message_count,
                "branched_from": chat_name,
                "target_floor": target_floor
            }
            
        except (FileExistsError, ValueError):
            raise
        except Exception as e:
            logger.error(f"创建分支失败 {role_name}/{chat_name}: {str(e)}")
            raise


# 全局实例
chat_service = ChatService(Path(settings.DATA_PATH))