# 会话管理器
# utils/session_manager.py
"""
会话管理器
支持连续对话，智能上下文管理，15万token限制
"""

import asyncio
import json
import os
import sys
from typing import Dict, List, Any, Optional
from pathlib import Path
from datetime import datetime
import requests
from dataclasses import dataclass
from aacode.i18n import t

TIKTOKEN_MIRRORS = [
    "https://raw.githubusercontent.com/openai/tiktoken/main/tiktoken/bpe/cl100k_base.tiktoken",
    "https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken",
    "https://raw.fastgit.org/openai/tiktoken/raw/main/tiktoken/bpe/cl100k_base.tiktoken",
    "https://gitclone.com/github.com/openai/tiktoken/raw/main/tiktoken/bpe/cl100k_base.tiktoken",
    "https://github.moeyy.xyz/https://raw.githubusercontent.com/openai/tiktoken/main/tiktoken/bpe/cl100k_base.tiktoken",
    "https://ghproxy.net/https://raw.githubusercontent.com/openai/tiktoken/main/tiktoken/bpe/cl100k_base.tiktoken",
    "https://mirror.ghproxy.com/https://raw.githubusercontent.com/openai/tiktoken/main/tiktoken/bpe/cl100k_base.tiktoken",
    "https://ghproxy.com/https://raw.githubusercontent.com/openai/tiktoken/main/tiktoken/bpe/cl100k_base.tiktoken",
]

TIKTOKEN_HASH = "223921b76ee99bde995b7ff738513eef100fb51d18c93597a113bcffe865b2a7"


def _download_tiktoken_file() -> Optional[str]:
    """尝试从多个镜像下载 tiktoken 文件"""
    import tempfile
    import hashlib

    cache_dir = os.path.join(tempfile.gettempdir(), "data-gym-cache")
    os.makedirs(cache_dir, exist_ok=True)
    cache_path = os.path.join(cache_dir, "cl100k_base.tiktoken")

    if os.path.exists(cache_path):
        with open(cache_path, "rb") as f:
            data = f.read()
        if hashlib.sha256(data).hexdigest() == TIKTOKEN_HASH:
            return cache_path

    for url in TIKTOKEN_MIRRORS:
        try:
            print(f"🔄 Attempting to download tiktoken from mirror: {url[:60]}...")
            resp = requests.get(url, timeout=10)
            resp.raise_for_status()
            data = resp.content

            actual_hash = hashlib.sha256(data).hexdigest()
            if actual_hash != TIKTOKEN_HASH:
                print(f"⚠️ Hash mismatch, skipping")
                continue

            with open(cache_path, "wb") as f:
                f.write(data)
            print(f"✅ tiktoken download successful: {cache_path}")
            return cache_path
        except Exception as e:
            print(f"❌ Download failed: {e}")
            continue

    return None


def _load_tiktoken_encoding():
    """加载 tiktoken 编码，支持镜像下载，添加超时保护"""
    import os
    import threading

    cache_file = os.path.expanduser("~/Library/Caches/tiktoken/cl100k_base.tiktoken")

    # 快速检查缓存文件是否存在且有效
    if os.path.exists(cache_file):
        file_size = os.path.getsize(cache_file)
        # 有效的 tiktoken 文件应该远大于 100 字节
        if file_size < 100:
            print(f"\n⚠️ tiktoken cache file corrupted (only {file_size} bytes), using simple estimation")
            return None

    result = [None]
    exception = [None]

    def _load():
        try:
            import tiktoken

            result[0] = tiktoken.get_encoding("cl100k_base")
        except Exception as e:
            exception[0] = e

    # 使用线程实现超时，避免主线程阻塞
    t = threading.Thread(target=_load, daemon=True)
    t.start()
    t.join(timeout=5)  # 5秒超时

    if t.is_alive():
        print(f"\n⚠️ tiktoken loading timeout, using simple estimation")
        return None

    if exception[0]:
        error_msg = str(exception[0])
        if (
            "Connection" in error_msg
            or "ConnectionResetError" in error_msg
            or "HTTPError" in error_msg
            or "404" in error_msg
            or "timed out" in error_msg.lower()
        ):
            print(f"\n⚠️ tiktoken download failed, using simple estimation")
            print(
                f"   Manual fix: https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken"
            )
            print(f"   Save to: ~/Library/Caches/tiktoken/cl100k_base.tiktoken")
        return None

    return result[0]


@dataclass
class SessionMessage:
    """会话消息"""

    role: str  # user, assistant, system, tool
    content: str
    timestamp: float
    tokens: int = 0
    metadata: Optional[Dict] = None
    tool_calls: Optional[List[Dict]] = None       # assistant 消息的 tool_calls
    tool_call_id: Optional[str] = None             # tool 消息的 tool_call_id
    reasoning_content: Optional[str] = None         # assistant 消息的 reasoning_content


@dataclass
class SessionSummary:
    """会话摘要"""

    session_id: str
    created_at: float
    last_activity: float
    total_messages: int
    total_tokens: int
    task_count: int
    title: str
    status: str  # active, completed, archived


class SessionManager:
    """会话管理器"""

    def __init__(self, project_path, max_tokens: int = 200000):
        from pathlib import Path
        self.project_path = Path(project_path) if not isinstance(project_path, Path) else project_path
        self.sessions_dir = self.project_path / ".aacode" / "sessions"
        self.sessions_dir.mkdir(parents=True, exist_ok=True)
        self.max_tokens = max_tokens

        # 当前会话
        self.current_session_id: Optional[str] = None
        self.current_messages: List[SessionMessage] = []
        self._structured_messages: Optional[List[Dict]] = None

        # 会话索引
        self.sessions_index: Dict[str, SessionSummary] = {}

        # token计数器
        self.encoding = _load_tiktoken_encoding()

        # 加载会话索引
        self._load_sessions_index()

    def _load_sessions_index(self):
        """加载会话索引"""
        index_file = self.sessions_dir / "sessions_index.json"
        if index_file.exists():
            try:
                with open(index_file, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    for session_id, session_data in data.items():
                        self.sessions_index[session_id] = SessionSummary(**session_data)
            except Exception as e:
                print(t("session.load_error", e=str(e)))

    def _save_sessions_index(self):
        """保存会话索引"""
        index_file = self.sessions_dir / "sessions_index.json"
        try:
            data = {}
            for session_id, session_summary in self.sessions_index.items():
                data[session_id] = {
                    "session_id": session_summary.session_id,
                    "created_at": session_summary.created_at,
                    "last_activity": session_summary.last_activity,
                    "total_messages": session_summary.total_messages,
                    "total_tokens": session_summary.total_tokens,
                    "task_count": session_summary.task_count,
                    "title": session_summary.title,
                    "status": session_summary.status,
                }

            with open(index_file, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
        except Exception as e:
            print(f"⚠️ Failed to save session index: {e}")

    def _count_tokens(self, text: str) -> int:
        """计算文本的token数量"""
        if self.encoding:
            try:
                return len(self.encoding.encode(text))
            except Exception:
                pass
        # 回退到简单估算（大致4字符=1token）
        return len(text) // 4

    def _get_total_tokens(self) -> int:
        """获取当前会话的总token数"""
        return sum(msg.tokens for msg in self.current_messages)

    async def create_session(self, task: str, title: Optional[str] = None) -> str:
        """创建新会话"""
        session_id = f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{len(self.sessions_index)}"

        # 生成标题
        if title is None:
            title = task[:50] + "..." if len(task) > 50 else task

        # 创建会话摘要
        session_summary = SessionSummary(
            session_id=session_id,
            created_at=datetime.now().timestamp(),
            last_activity=datetime.now().timestamp(),
            total_messages=0,
            total_tokens=0,
            task_count=0,
            title=title,
            status="active",
        )

        self.sessions_index[session_id] = session_summary
        self.current_session_id = session_id
        self.current_messages = []

        # 添加系统消息和任务
        system_msg = SessionMessage(
            role="system",
            content=f"You are an AI coding assistant. Please work as follows:\n1. First plan the task steps\n2. Execute and verify step by step\n3. Ensure code quality and maintainability",
            timestamp=datetime.now().timestamp(),
            tokens=self._count_tokens(
                f"You are an AI coding assistant. Please work as follows:\n1. First plan the task steps\n2. Execute and verify step by step\n3. Ensure code quality and maintainability"
            ),
        )

        self.current_messages = [system_msg]
        self._structured_messages = None
        # 只有 task 非空时才添加 user 消息
        if task and task.strip():
            user_msg = SessionMessage(
                role="user",
                content=task,
                timestamp=datetime.now().timestamp(),
                tokens=self._count_tokens(task),
            )
            self.current_messages.append(user_msg)
        session_summary.total_messages = len(self.current_messages)
        session_summary.total_tokens = self._get_total_tokens()

        # 保存索引
        self._save_sessions_index()

        # 保存会话文件
        await self._save_session()

        # 保存当前会话ID到文件（方便用户查询）
        await self._save_current_session_id()

        return session_id

    async def _save_current_session_id(self):
        """保存当前会话ID到文件"""
        if not self.current_session_id:
            return

        current_session_file = self.sessions_dir.parent / "current_session.txt"
        try:
            with open(current_session_file, "w", encoding="utf-8") as f:
                f.write(f"Session ID: {self.current_session_id}\n")
                f.write(
                    f"Created: {datetime.fromtimestamp(self.sessions_index[self.current_session_id].created_at).strftime('%Y-%m-%d %H:%M:%S')}\n"
                )
                f.write(f"Title: {self.sessions_index[self.current_session_id].title}\n")
                f.write(f"Messages: {len(self.current_messages)}\n")
                f.write(f"Tokens: {self._get_total_tokens()}\n")
        except Exception as e:
            print(t("session.save_id_error", e=str(e)))

    async def add_message(
        self, role: str, content: str, metadata: Optional[Dict] = None,
        tool_calls: Optional[List[Dict]] = None,
        tool_call_id: Optional[str] = None,
        reasoning_content: Optional[str] = None,
    ) -> bool:
        """添加消息到当前会话"""
        if not self.current_session_id:
            return False

        # 检查token限制
        new_tokens = self._count_tokens(content)
        current_tokens = self._get_total_tokens()

        if current_tokens + new_tokens > self.max_tokens:
            await self._compact_context()
            current_tokens = self._get_total_tokens()
            if current_tokens + new_tokens > self.max_tokens:
                return False

        # 添加消息
        message = SessionMessage(
            role=role,
            content=content,
            timestamp=datetime.now().timestamp(),
            tokens=new_tokens,
            metadata=metadata or {},
            tool_calls=tool_calls,
            tool_call_id=tool_call_id,
            reasoning_content=reasoning_content,
        )

        self.current_messages.append(message)

        # 如果 _structured_messages 已存在，同步追加（保持双源一致）
        if self._structured_messages is not None:
            entry = {
                "role": role,
                "content": content,
                "tool_calls": tool_calls,
                "tool_call_id": tool_call_id,
                "reasoning_content": reasoning_content,
            }
            # 移除 None 值，保持 JSON 整洁
            entry = {k: v for k, v in entry.items() if v is not None}
            self._structured_messages.append(entry)

        # 更新会话摘要
        if self.current_session_id in self.sessions_index:
            session_summary = self.sessions_index[self.current_session_id]
            session_summary.last_activity = datetime.now().timestamp()
            session_summary.total_messages = len(self.current_messages)
            session_summary.total_tokens = self._get_total_tokens()
            # 如果 title 为空，用第一条 user 消息作为标题
            if role == "user" and (not session_summary.title or session_summary.title.strip() == ""):
                session_summary.title = content[:50] + ("..." if len(content) > 50 else "")

        # 保存
        await self._save_session()
        self._save_sessions_index()

        # 更新当前会话ID文件
        await self._save_current_session_id()

        return True

    async def _save_structured_messages(self, messages: List[Dict]):
        """保存结构化消息到会话文件（含 tool_calls / reasoning_content），与文本消息并存"""
        if not self.current_session_id:
            return
        session_file = self.sessions_dir / f"{self.current_session_id}.json"
        if not session_file.exists():
            return
        try:
            with open(session_file, "r", encoding="utf-8") as f:
                session_data = json.load(f)
            session_data["structured_messages"] = messages
            with open(session_file, "w", encoding="utf-8") as f:
                json.dump(session_data, f, indent=2, ensure_ascii=False)
            self._structured_messages = messages
        except Exception as e:
            print(f"⚠️ Failed to save structured message: {e}")

    async def get_messages(
        self, session_id: Optional[str] = None, include_system: bool = True
    ) -> List[Dict]:
        """获取会话消息"""
        target_session_id = session_id or self.current_session_id
        if not target_session_id:
            return []

        # 如果不是当前会话，加载它
        if target_session_id != self.current_session_id:
            await self._load_session(target_session_id)

        # 优先返回结构化消息（含 tool_calls / reasoning_content）
        if self._structured_messages:
            filtered = []
            for msg in self._structured_messages:
                if not include_system and msg.get("role") == "system":
                    continue
                filtered.append(msg)
            return filtered

        messages = []
        for msg in self.current_messages:
            if not include_system and msg.role == "system":
                continue
            entry = {"role": msg.role, "content": msg.content}
            if msg.tool_calls:
                entry["tool_calls"] = msg.tool_calls
            if msg.tool_call_id:
                entry["tool_call_id"] = msg.tool_call_id
            if msg.reasoning_content:
                entry["reasoning_content"] = msg.reasoning_content
            messages.append(entry)

        return messages

    async def get_conversation_history(self, max_length: int = 10) -> str:
        """获取对话历史（用于显示给用户）"""
        if not self.current_messages:
            return ""

        recent_messages = self.current_messages[-max_length:]
        history_lines = []

        for msg in recent_messages:
            if msg.role == "system":
                continue
            role_name = "User" if msg.role == "user" else "Assistant"
            content_preview = (
                msg.content[:100] + "..." if len(msg.content) > 100 else msg.content
            )
            history_lines.append(f"{role_name}: {content_preview}")

        return "\n".join(history_lines)

    async def _compact_context(self):
        """压缩上下文，保持在token限制内"""
        if len(self.current_messages) <= 4:  # 保留系统消息和最少对话
            return

        # 保留系统消息和最近的用户消息
        system_msgs = [msg for msg in self.current_messages if msg.role == "system"]
        recent_msgs = self.current_messages[-3:]  # 保留最近3条消息

        # 生成摘要
        old_msgs = self.current_messages[:-3]
        if old_msgs:
            compact_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            summary = f"The previous conversation contained {len(old_msgs)} messages, mainly discussing programming-related tasks."
            summary_msg = SessionMessage(
                role="system",
                content=f"Context summary (compressed at: {compact_time}): {summary}",
                timestamp=datetime.now().timestamp(),
                tokens=self._count_tokens(summary),
            )

            self.current_messages = system_msgs + [summary_msg] + recent_msgs
        else:
            self.current_messages = system_msgs + recent_msgs

    async def _save_session(self):
        """保存当前会话（保留文件中已有的 structured_messages 等字段）"""
        if not self.current_session_id:
            return

        session_file = self.sessions_dir / f"{self.current_session_id}.json"

        # 先读取已有文件，保留 structured_messages 等非冲突字段
        existing = {}
        if session_file.exists():
            try:
                with open(session_file, "r", encoding="utf-8") as f:
                    existing = json.load(f)
            except Exception:
                existing = {}

        session_data = {
            "session_id": self.current_session_id,
            "created_at": self.sessions_index.get(
                self.current_session_id, SessionSummary("", 0, 0, 0, 0, 0, "", "")
            ).created_at,
            "messages": [
                {
                    "role": msg.role,
                    "content": msg.content,
                    "timestamp": msg.timestamp,
                    "tokens": msg.tokens,
                    "metadata": msg.metadata,
                    "tool_calls": msg.tool_calls,
                    "tool_call_id": msg.tool_call_id,
                    "reasoning_content": msg.reasoning_content,
                }
                for msg in self.current_messages
            ],
        }

        # 优先用内存中的 _structured_messages（更实时），否则保留文件中已有的
        if self._structured_messages is not None:
            session_data["structured_messages"] = self._structured_messages
        elif "structured_messages" in existing:
            session_data["structured_messages"] = existing["structured_messages"]

        try:
            with open(session_file, "w", encoding="utf-8") as f:
                json.dump(session_data, f, indent=2, ensure_ascii=False)
        except Exception as e:
            print(t("session.save_error", e=str(e)))

    async def _load_session(self, session_id: str):
        """加载指定会话"""
        session_file = self.sessions_dir / f"{session_id}.json"
        if not session_file.exists():
            return False

        try:
            with open(session_file, "r", encoding="utf-8") as f:
                session_data = json.load(f)

            self.current_session_id = session_id
            self.current_messages = []
            self._structured_messages: Optional[List[Dict]] = session_data.get("structured_messages")

            for msg_data in session_data.get("messages", []):
                message = SessionMessage(
                    role=msg_data["role"],
                    content=msg_data["content"],
                    timestamp=msg_data["timestamp"],
                    tokens=msg_data.get("tokens", 0),
                    metadata=msg_data.get("metadata"),
                    tool_calls=msg_data.get("tool_calls"),
                    tool_call_id=msg_data.get("tool_call_id"),
                    reasoning_content=msg_data.get("reasoning_content"),
                )
                self.current_messages.append(message)

            return True

        except Exception as e:
            print(t("session.load_session_error", e=str(e)))
            return False

    async def list_sessions(self) -> List[Dict]:
        """列出所有会话"""
        sessions = []
        for session_summary in self.sessions_index.values():
            sessions.append(
                {
                    "session_id": session_summary.session_id,
                    "title": session_summary.title,
                    "created_at": session_summary.created_at,
                    "last_activity": session_summary.last_activity,
                    "total_messages": session_summary.total_messages,
                    "total_tokens": session_summary.total_tokens,
                    "task_count": session_summary.task_count,
                    "status": session_summary.status,
                }
            )

        # 按最后活动时间排序（处理None值）
        sessions.sort(key=lambda x: float(x.get("last_activity", 0) or 0), reverse=True)
        return sessions

    async def switch_session(self, session_id: str) -> bool:
        """切换到指定会话"""
        if session_id not in self.sessions_index:
            return False

        # 保存当前会话
        if self.current_session_id:
            await self._save_session()

        # 加载目标会话
        success = await self._load_session(session_id)
        if success:
            self.current_session_id = session_id
            # 更新当前会话ID文件
            await self._save_current_session_id()

        return success

    async def delete_session(self, session_id: str) -> bool:
        """删除会话"""
        if session_id not in self.sessions_index:
            return False

        # 删除会话文件
        session_file = self.sessions_dir / f"{session_id}.json"
        if session_file.exists():
            session_file.unlink()

        # 从索引中删除
        del self.sessions_index[session_id]
        self._save_sessions_index()

        # 如果是当前会话，清空
        if self.current_session_id == session_id:
            self.current_session_id = None
            self.current_messages = []

        return True

    async def archive_session(self, session_id: str) -> bool:
        """归档会话"""
        if session_id not in self.sessions_index:
            return False

        # 更新状态
        self.sessions_index[session_id].status = "archived"
        self._save_sessions_index()

        return True

    async def get_session(self, session_id: Optional[str] = None) -> Dict[str, Any]:
        """获取会话详细信息"""
        if not session_id:
            session_id = self.current_session_id

        if not session_id:
            return {}

        # 加载会话数据
        session_file = self.sessions_dir / f"{session_id}.json"
        if not session_file.exists():
            return {}

        try:
            with open(session_file, "r", encoding="utf-8") as f:
                session_summary = json.load(f)
        except Exception:
            return {}

        if not session_summary:
            return {}

        messages = await self.get_messages(session_id, include_system=True)

        return {
            "session_id": session_id,
            "title": session_summary.get("title", ""),
            "messages": messages,
            "total_messages": session_summary.get("total_messages", 0),
            "total_tokens": session_summary.get("total_tokens", 0),
            "created_at": session_summary.get("created_at", ""),
            "last_activity": session_summary.get("last_activity", ""),
        }

    async def count_tokens(self, session_id: Optional[str] = None) -> int:
        """计算会话的token数量"""
        if not session_id:
            session_id = self.current_session_id

        if not session_id:
            return 0

        # 加载会话数据
        session_file = self.sessions_dir / f"{session_id}.json"
        if not session_file.exists():
            return 0

        try:
            with open(session_file, "r", encoding="utf-8") as f:
                session_data = json.load(f)
            return session_data.get("total_tokens", 0)
        except Exception:
            return 0

    @property
    def token_limit(self) -> int:
        """获取token限制"""
        return self.max_tokens

    async def compress_session(
        self, session_id: Optional[str] = None, max_messages: int = 50
    ) -> Dict[str, Any]:
        """压缩会话消息"""
        if not session_id:
            session_id = self.current_session_id

        if not session_id:
            return {"compressed_messages": [], "removed_count": 0}

        original_messages = await self.get_messages(session_id, include_system=False)

        if len(original_messages) <= max_messages:
            return {"compressed_messages": original_messages, "removed_count": 0}

        # 保留最近的max_messages条消息
        compressed_messages = original_messages[-max_messages:]
        removed_count = len(original_messages) - max_messages

        return {
            "compressed_messages": compressed_messages,
            "removed_count": removed_count,
        }

    def get_session_stats(self) -> Dict[str, Any]:
        """获取当前会话统计信息"""
        if not self.current_session_id:
            return {}

        # 加载会话数据
        session_file = self.sessions_dir / f"{self.current_session_id}.json"
        if not session_file.exists():
            return {}

        try:
            with open(session_file, "r", encoding="utf-8") as f:
                session_summary = json.load(f)
        except Exception:
            return {}

        return {
            "session_id": self.current_session_id,
            "title": session_summary.get("title", ""),
            "total_messages": session_summary.get("total_messages", 0),
            "total_tokens": session_summary.get("total_tokens", 0),
            "current_tokens": self._get_total_tokens(),
            "max_tokens": self.max_tokens,
            "created_at": session_summary.get("created_at", ""),
            "last_activity": session_summary.get("last_activity", ""),
        }
