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

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

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 {"chat": [{"role_name": role, **chat} for role, chats in result.items() for chat in chats]}
    
    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 {
                "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 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)
        
        # 构建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
        }
        
        # 写入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