"""
提示词组装器 (Prompt Assembler)

负责根据 SillyTavern 规范将角色卡、世界书、聊天历史等组件
拼装成最终的 LLM 消息列表。
"""
import re
from typing import List, Dict, Optional
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, BaseMessage

from models.internal import CharacterCard, ChatMessage, WorldInfoEntry


class PromptConfig:
    """提示词组装配置"""
    def __init__(
        self,
        an_position: str = "after_history", # "before_history" or "after_history"
        an_depth: int = 4,
        post_history_instructions: Optional[str] = None
    ):
        self.an_position = an_position
        self.an_depth = an_depth
        self.post_history_instructions = post_history_instructions


class PromptAssembler:
    """
    轻量级提示词组装核心
    
    不依赖复杂的框架，只负责纯粹的文本拼接和位置插入。
    """

    # SillyTavern 的位置枚举映射
    POS_WI_BEFORE = 0
    POS_WI_AFTER = 1
    POS_EXAMPLES_BEFORE = 2
    POS_EXAMPLES_AFTER = 3
    POS_AN_TOP = 4
    POS_AN_BOTTOM = 5
    POS_DEPTH = 6
    POS_OUTLET = 7

    def assemble(
        self,
        character: CharacterCard,
        chat_history: List[ChatMessage],
        user_input: str,
        active_entries: List[WorldInfoEntry],
        config: PromptConfig = PromptConfig()
    ) -> List[BaseMessage]:
        """
        执行完整的提示词组装流程
        
        Returns:
            List[BaseMessage]: 准备好发送给 LLM 的消息列表
        """
        # 1. 按位置分组世界书条目
        grouped_entries = self._group_entries_by_position(active_entries)

        # 2. 组装 Story String (包含 Pos 0-3)
        story_string = self._build_story_string(character, grouped_entries)

        # 3. 组装 Author's Note (包含 Pos 4-5)
        authors_note_content = self._build_authors_note(grouped_entries, config.an_depth)

        # 4. 处理 Chat History 并注入 Depth 条目 (Pos 6)
        processed_history = self._inject_depth_entries(chat_history, grouped_entries.get(self.POS_DEPTH, []))

        # 5. 准备 Outlet 替换字典 (Pos 7)
        outlet_map = {entry.uid: entry.content for entry in grouped_entries.get(self.POS_OUTLET, [])}

        # 6. 最终封装为 Messages
        return self._wrap_to_messages(
            story_string, 
            authors_note_content, 
            processed_history, 
            user_input, 
            outlet_map,
            config
        )

    def _group_entries_by_position(self, entries: List[WorldInfoEntry]) -> Dict[int, List[WorldInfoEntry]]:
        """将激活的条目按 position 分组"""
        grouped = {}
        for entry in entries:
            # 这里假设 entry.position 存储的是我们定义的 0-7 整数
            pos = entry.position if isinstance(entry.position, int) else 1 # 默认为 wiAfter
            if pos not in grouped:
                grouped[pos] = []
            grouped[pos].append(entry)
        
        # 对每个组内的条目按 order 排序
        for pos in grouped:
            grouped[pos].sort(key=lambda x: x.order)
        return grouped

    def _build_story_string(self, character: CharacterCard, grouped: Dict) -> str:
        """组装故事字符串 (Story String)"""
        parts = []
        
        # Pos 0: wiBefore
        for entry in grouped.get(self.POS_WI_BEFORE, []):
            parts.append(entry.content)
            
        # 角色核心信息
        parts.append(f"[Character('{character.name}')]\n{character.description}\n")
        parts.append(f"Personality: {character.personality}\n")
        parts.append(f"Scenario: {character.scenario}\n")
        
        # Pos 1: wiAfter
        for entry in grouped.get(self.POS_WI_AFTER, []):
            parts.append(entry.content)
            
        # Pos 2: Examples Before
        for entry in grouped.get(self.POS_EXAMPLES_BEFORE, []):
            parts.append(entry.content)
            
        # 示例对话
        if character.mes_example:
            parts.append(f"<START>\n{character.mes_example}")
            
        # Pos 3: Examples After
        for entry in grouped.get(self.POS_EXAMPLES_AFTER, []):
            parts.append(entry.content)
            
        return "\n".join(parts)

    def _build_authors_note(self, grouped: Dict, depth: int) -> str:
        """组装作者笔记 (Author's Note)"""
        parts = []
        
        # Pos 4: AN Top
        for entry in grouped.get(self.POS_AN_TOP, []):
            parts.append(str(entry.content) if entry.content else "")
            
        # AN 核心内容 (这里简化为一个占位，实际应从角色卡或设置获取)
        parts.append(f"[Author's note at depth {depth}]")
        
        # Pos 5: AN Bottom
        for entry in grouped.get(self.POS_AN_BOTTOM, []):
            parts.append(str(entry.content) if entry.content else "")
            
        return "\n".join(parts)

    def _inject_depth_entries(self, history: List[ChatMessage], depth_entries: List[WorldInfoEntry]) -> List[Dict]:
        """
        在聊天历史的指定深度插入条目 (Pos 6)
        返回一个包含 role 和 content 的字典列表，方便后续转换
        
        ✅ 过滤已被总结的消息（is_summarized=True 且 mes=""）
        """
        # 先将历史转换为中间格式，过滤掉空消息（已被总结）
        msg_list = []
        for msg in history:
            # ✅ 跳过已被总结的空消息
            if msg.is_summarized and (msg.mes == "" or msg.mes.strip() == ""):
                continue
            msg_list.append({"role": "user" if msg.is_user else "assistant", "content": msg.mes})
        
        # 按 depth 分组插入
        # d0 通常指最新用户输入之前，即列表末尾
        for entry in depth_entries:
            depth = entry.depth if entry.depth is not None else 0
            # 计算插入索引 (从后往前数)
            insert_index = max(0, len(msg_list) - depth)
            
            # 确定角色
            role_map = {"system": "system", "user": "user", "assistant": "assistant"}
            role = role_map.get(str(entry.position).split('_')[-1] if '_' in str(entry.position) else "system", "system")
            
            msg_list.insert(insert_index, {"role": "system", "content": entry.content})
            
        return msg_list

    def _replace_outlets(self, text: str, outlet_map: Dict[str, str]) -> str:
        """执行 Outlet 宏替换 (Pos 7)"""
        def replacer(match):
            uid = match.group(1)
            return outlet_map.get(uid, "")
        
        # 匹配 {{outlet::UID}}
        return re.sub(r"\{\{outlet::([^}]+)\}\}", replacer, text)

    def _wrap_to_messages(
        self, 
        story_string: str, 
        an_content: str, 
        history: List[Dict], 
        user_input: str,
        outlet_map: Dict[str, str],
        config: PromptConfig
    ) -> List[BaseMessage]:
        """将组装好的文本块封装为 LangChain Messages"""
        messages = []
        
        # 1. System Message (Story String + Outlet 替换)
        final_story = self._replace_outlets(story_string, outlet_map)
        if final_story:
            messages.append(SystemMessage(content=final_story))
            
        # 2. Author's Note (根据配置位置插入)
        if an_content and config.an_position == "before_history":
            messages.append(SystemMessage(content=self._replace_outlets(an_content, outlet_map)))
            
        # 3. Chat History
        for msg_data in history:
            if msg_data["role"] == "user":
                messages.append(HumanMessage(content=msg_data["content"]))
            elif msg_data["role"] == "assistant":
                messages.append(AIMessage(content=msg_data["content"]))
            else:
                messages.append(SystemMessage(content=msg_data["content"]))
                
        # 4. Author's Note (如果在 History 之后)
        if an_content and config.an_position == "after_history":
            messages.append(SystemMessage(content=self._replace_outlets(an_content, outlet_map)))
            
        # 5. Post-History Instructions & User Input
        final_input = user_input
        if config.post_history_instructions:
            final_input = f"{config.post_history_instructions}\n\n{user_input}"
            
        messages.append(HumanMessage(content=final_input))
        
        return messages
