# -*- coding: utf-8 -*-
"""
会话管理
管理一次对话会话的消息历史，支持上下文维护和持久化恢复。
"""
import json
import os
import time
from pathlib import Path
from typing import List, Dict, Any


class Session:
    """管理一次对话会话的消息历史，维护完整的对话上下文。
    支持 save() 持久化和 from_file() 恢复。
    """

    def __init__(self, system_prompt: str):
        self.messages: List[Dict[str, Any]] = [
            {"role": "system", "content": system_prompt}
        ]
        self._tool_call_counter = 0
        # 记录会话启动时间戳（/stats 命令使用）
        self.started_at = time.time()
        # 记录会话工作目录（恢复会话时切换回去，避免 PathSecurityService 边界错位）
        self.cwd = os.getcwd()

    def save(self, path):
        """保存会话到 JSON 文件。原子写入：先写 .tmp 再 replace。"""
        data = {
            "messages": self.messages,
            "tool_call_counter": self._tool_call_counter,
            "started_at": getattr(self, "started_at", None),
            "cwd": getattr(self, "cwd", None),
        }
        p = Path(path)
        p.parent.mkdir(parents=True, exist_ok=True)
        tmp = p.with_suffix(".tmp")
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False)
        tmp.replace(p)

    @classmethod
    def from_file(cls, path, system_prompt: str) -> "Session":
        """从 JSON 文件恢复会话。文件不存在或损坏时返回全新 session。

        会话恢复时主动清洗历史：截断尾部不完整的 tool_calls 序列，
        避免 assistant(tool_calls) 后缺 tool 响应导致 API 400 错误。
        场景：用户按 Esc 中断任务后保存会话，最后一条 assistant 的
        tool_calls 可能没有对应的 tool 角色响应。
        """
        p = Path(path)
        if not p.exists():
            return cls(system_prompt)
        try:
            with open(p, "r", encoding="utf-8") as f:
                data = json.load(f)
            session = cls.__new__(cls)
            messages = data.get("messages", [
                {"role": "system", "content": system_prompt}
            ])
            session.messages = cls._sanitize_messages(messages, system_prompt)
            session._tool_call_counter = data.get("tool_call_counter", 0)
            # 恢复会话启动时间戳，缺失时用当前时间（避免 /stats 显示「未知」）
            session.started_at = data.get("started_at", time.time())
            # 恢复会话工作目录：若原会话记录了 cwd 且目录仍存在，切换回去，
            # 避免 PathSecurityService 边界错位导致 file_write 误判。
            # 切换失败时静默降级（保留当前 CWD），不阻断会话恢复。
            saved_cwd = data.get("cwd")
            if saved_cwd and os.path.isdir(saved_cwd):
                try:
                    os.chdir(saved_cwd)
                    session.cwd = saved_cwd
                except (OSError, ValueError):
                    session.cwd = os.getcwd()
            else:
                session.cwd = os.getcwd()
            return session
        except (json.JSONDecodeError, KeyError, OSError):
            return cls(system_prompt)

    @staticmethod
    def _sanitize_messages(
        messages: List[Dict[str, Any]],
        system_prompt: str = None,
    ) -> List[Dict[str, Any]]:
        """清洗历史消息序列，确保满足 API 协议。

        API 协议要求：每条 assistant(tool_calls) 后必须紧跟覆盖所有
        tool_call_id 的 tool 角色消息。若发现不完整序列，从该 assistant
        处截断后续所有消息（避免悬挂的 user/assistant）。

        同时保证 messages[0] 是 system：若被截断或本就缺失，且
        system_prompt 非空，则在首位补一条 system 消息。

        返回新的列表，不修改输入。
        """
        if not messages:
            if system_prompt:
                return [{"role": "system", "content": system_prompt}]
            return []

        # 找到第一条不合法 assistant 的索引；找不到则全部合法
        cut_index = len(messages)
        for i, msg in enumerate(messages):
            if msg.get("role") != "assistant":
                continue
            tool_calls = msg.get("tool_calls")
            if not tool_calls:
                continue
            # 收集紧随其后连续 tool 消息的 tool_call_id
            expected_ids = {tc.get("id") for tc in tool_calls if tc.get("id")}
            got_ids = set()
            j = i + 1
            while j < len(messages) and messages[j].get("role") == "tool":
                tc_id = messages[j].get("tool_call_id")
                if tc_id:
                    got_ids.add(tc_id)
                j += 1
            # 若有未配对的 tool_call_id，从该 assistant 处截断
            if not expected_ids.issubset(got_ids):
                cut_index = i
                break

        result = messages[:cut_index]

        # 保证 messages[0] 是 system（防御上下文压缩丢失 system）
        if not result:
            if system_prompt:
                result = [{"role": "system", "content": system_prompt}]
        elif result[0].get("role") != "system":
            if system_prompt:
                result = [{"role": "system", "content": system_prompt}] + result
            else:
                # 没有 system_prompt 兜底，从 messages 中找已有的 system
                sys_msg = next(
                    (m for m in result if m.get("role") == "system"), None
                )
                if sys_msg:
                    result = [sys_msg] + [m for m in result if m is not sys_msg]

        return result

    @staticmethod
    def _ensure_system_preserved(
        messages: List[Dict[str, Any]],
        system_prompt: str,
    ) -> List[Dict[str, Any]]:
        """守卫：供 main.py 上下文压缩后调用，保证 messages[0] 是 system。

        场景：上下文压缩引擎替换 session.messages 时可能返回不含 system
        的列表，导致 LLM 调用丢失系统指令。此方法在替换前调用，
        若 system 缺失则在首位补一条。
        """
        if not messages:
            return [{"role": "system", "content": system_prompt}]
        if messages[0].get("role") == "system":
            return messages
        return [{"role": "system", "content": system_prompt}] + messages

    def add_user(self, content: str):
        """添加用户消息"""
        self.messages.append({"role": "user", "content": content})

    def append_to_system(self, text: str) -> None:
        """把文本追加到 system 消息末尾。供 LifecycleHost 注入 skill 内容用。"""
        if self.messages and self.messages[0]["role"] == "system":
            self.messages[0]["content"] += text

    def add_assistant(self, content: str = None, tool_calls: list = None, reasoning_content: str = None):
        """添加助手消息（含工具调用和思考链）。

        不变式：assistant 消息必须满足 API 协议要求 —— content 或 tool_calls 至少有一个。
        若调用方只传 reasoning_content（模型只思考不出内容/工具），自动补 content=""
        以保证 get_messages_for_api() 剔除 reasoning_content 后消息仍合法，不会触发 400。
        """
        msg = {"role": "assistant"}
        if content is not None:
            msg["content"] = content
        if reasoning_content:
            msg["reasoning_content"] = reasoning_content
        if tool_calls:
            msg["tool_calls"] = tool_calls
        if "content" not in msg and "tool_calls" not in msg:
            msg["content"] = ""
        self.messages.append(msg)

    def rollback_last_assistant(self):
        """回滚最后一条 assistant 消息（含 tool_calls 但无对应 tool 结果时使用）。

        场景：agent_loop.process_tool_calls 在 add_assistant(tool_calls=...) 之后
        工具执行被 TaskInterrupted 中断，若不回滚会留下孤儿消息，导致下一次
        API 调用报 400 "insufficient tool messages following tool_calls"。
        仅当末尾确实是 assistant 且后面没有 tool 消息时才删除。
        """
        if (len(self.messages) >= 1
                and self.messages[-1].get("role") == "assistant"):
            self.messages.pop()

    def add_tool_result(self, tool_call_id: str, content: str, max_chars: int = None):
        """添加工具执行结果。max_chars 非空时截断过长内容。"""
        # ponytail: 截断过长的工具结果，避免上下文爆炸
        if max_chars and len(content) > max_chars:
            content = content[:max_chars] + f"\n...[截断, 共 {len(content)} 字符]"
        self.messages.append({
            "role": "tool",
            "tool_call_id": tool_call_id,
            "content": content
        })
        self._tool_call_counter += 1

    def auto_name(self) -> str:
        """从首条用户消息自动生成会话名。取前 30 字符，只保留中英文数字。"""
        import re
        from datetime import datetime
        for msg in self.messages:
            if msg["role"] == "user":
                text = re.sub(r'[^\w\u4e00-\u9fff]', '', msg.get("content", ""))[:30]
                return f"{text}_{datetime.now():%m%d-%H%M}" if text else f"session_{datetime.now():%m%d-%H%M}"
        return f"session_{datetime.now():%m%d-%H%M}"

    @staticmethod
    def list_sessions(session_dir: Path) -> list:
        """列出所有已保存的会话名称（latest 排最前）。"""
        if not session_dir.exists():
            return []
        return sorted(
            [p.stem for p in session_dir.glob("*.json")],
            key=lambda n: (n != "latest", n)
        )

    @staticmethod
    def delete_session(session_dir: Path, name: str):
        """删除指定会话文件。"""
        (session_dir / f"{name}.json").unlink(missing_ok=True)

    def get_messages(self) -> List[Dict[str, Any]]:
        """获取当前消息列表（用于 API 请求）"""
        return self.messages

    def get_messages_for_api(self) -> List[Dict[str, Any]]:
        """获取发给 LLM API 的消息列表。剔除 reasoning_content 以节省 token。

        reasoning_content 是 AI 自己的上次思考，AI 不需要回看旧思考。
        协议设计参考 OpenAI/Claude 协议，reasoning_content 是单次响应的辅助字段。
        本地审计 get_messages() 仍保留 reasoning_content。

        出口兜底校验：调用 _sanitize_messages 截断任何孤儿 assistant(tool_calls)
        消息，防止上一轮工具执行被中断后残留的非法序列触发 API 400。
        """
        sanitized = self._sanitize_messages(self.messages)
        if len(sanitized) < len(self.messages):
            # 检测到孤儿，同步清理 session.messages 避免后续重复触发
            self.messages = sanitized
        result = []
        for m in sanitized:
            if "reasoning_content" in m:
                m = {k: v for k, v in m.items() if k != "reasoning_content"}
            result.append(m)
        return result

    def get_message_count(self) -> int:
        """获取消息数量（不含 system prompt）"""
        return len(self.messages) - 1

    def get_tool_call_count(self) -> int:
        """获取工具调用次数"""
        return self._tool_call_counter

    def clear(self):
        """清空历史，保留 system prompt"""
        system = self.messages[0]
        self.messages = [system]
        self._tool_call_counter = 0

    def get_last_user_message(self) -> Dict[str, Any]:
        """获取最后一条用户消息"""
        for msg in reversed(self.messages):
            if msg["role"] == "user":
                return msg
        return None