"""
默认 Context 实现

特性：
- 两阶段压缩（Prune → LLM 摘要）
- 优先使用 LLM 返回的实际 token 数
- 可选持久化：每次 add_message 异步写盘（JSONL 格式）
- 支持从磁盘恢复上下文
"""

import asyncio
import json
from typing import List, Optional
from pathlib import Path

from provider import (
    Message,
    MessageRole,
    ToolResult,
    TokenUsage,
)
from core.logger import log_event, log_error

from .base import BaseContext, ContextConfig, ChatFn, CountTokensFn


def _load_compaction_prompt() -> str:
    """加载压缩提示词"""
    prompt_path = Path(__file__).parent.parent / "assets" / "prompts" / "utilities" / "compaction.txt"
    if prompt_path.exists():
        return prompt_path.read_text(encoding="utf-8")
    return "请总结以下对话历史，保留关键信息，移除冗余内容。"


class DefaultContext(BaseContext):
    """
    默认上下文实现

    压缩策略：
    1. Prune: 清理旧的工具输出（替换为占位符）
    2. Summarize: 使用 LLM 生成摘要（仅当 Prune 不够时）

    持久化：
    - session_id 非空时启用，JSONL 格式追加写入
    - 压缩后重写整个文件
    - 支持 load_from_disk() 恢复
    """

    TOOL_OUTPUT_PLACEHOLDER = "[旧工具输出已清理]"

    def __init__(
        self,
        chat_fn: ChatFn,
        config: Optional[ContextConfig] = None,
        session_id: Optional[str] = None,
        storage_dir: Optional[Path] = None,
        count_tokens_fn: Optional[CountTokensFn] = None,
    ):
        super().__init__(chat_fn, config)

        self._system_prompts: List[Message] = []
        self._messages: List[Message] = []
        self._token_count: int = 0

        # 压缩锁
        self._compressing: bool = False

        # 压缩提示词
        self._compaction_prompt = _load_compaction_prompt()

        # Token 精确计数回调（可选，由 Agent 传入）
        self._count_tokens_fn = count_tokens_fn

        # 持久化（session_id 非空时启用）
        self._context_file: Optional[Path] = None
        self._write_lock = asyncio.Lock()

        if session_id:
            if not storage_dir:
                from core.paths import get_path_manager
                storage_dir = get_path_manager().get_data_dir() / "sessions"
            self._context_file = storage_dir / session_id / "context.jsonl"
            self._context_file.parent.mkdir(parents=True, exist_ok=True)

    # ==================== 公开方法 ====================

    def add_message(self, message: Message) -> None:
        """添加消息，异步写盘"""
        self._messages.append(message)

        if self._context_file:
            try:
                asyncio.get_running_loop().create_task(self._append_to_disk(message))
            except RuntimeError:
                pass  # 无事件循环时跳过

    @property
    def needs_compress(self) -> bool:
        """是否需要压缩"""
        return self._check_overflow() and not self._compressing

    async def compress_if_needed(self, force: bool = False) -> bool:
        """如果需要压缩则执行（由 agent 主循环显式调用）

        Args:
            force: 是否强制执行压缩（忽略阈值检查）
        """
        # 强制模式下跳过阈值检查
        if not force and not self.needs_compress:
            return False
        self._compressing = True
        await self._compress_async()
        return True

    def get_messages(self) -> List[Message]:
        """获取对话消息列表（不含系统提示词）"""
        return list(self._messages)

    def update_token_usage(self, usage: Optional[TokenUsage]) -> None:
        """更新 token 使用量"""
        if not usage:
            return

        # 使用 context_tokens（input + cache_creation + cache_read）反映真实上下文大小
        context_tokens = usage.total_tokens
        if context_tokens > 0:
            self._token_count = context_tokens

            log_event(
                "context_token_update",
                input_tokens=usage.input_tokens,
                cache_creation_input_tokens=usage.cache_creation_input_tokens,
                cache_read_input_tokens=usage.cache_read_input_tokens,
                context_tokens=context_tokens,
                max_tokens=self.config.max_tokens,
                usage_percent=f"{(context_tokens / self.config.max_tokens * 100):.1f}%",
            )

    def set_system_prompt(self, prompts: List[Message]) -> None:
        """设置系统提示词"""
        self._system_prompts = prompts

    @property
    def system_prompts(self) -> List[Message]:
        """获取系统提示词列表"""
        return self._system_prompts

    def clear(self) -> None:
        """清空上下文（保留系统提示词），同时清理磁盘文件"""
        self._messages.clear()
        self._token_count = 0

        if self._context_file and self._context_file.exists():
            try:
                self._context_file.unlink()
            except Exception as e:
                log_error(f"context_clear_disk_error: {str(e)}")

    def get_token_count(self) -> int:
        """获取当前 token 使用量"""
        if self._token_count > 0:
            return self._token_count
        return self._estimate_tokens()

    # ==================== 内部方法 ====================

    def _check_overflow(self, threshold: float = None) -> bool:
        """
        检查上下文是否溢出

        Args:
            threshold: 溢出阈值（0.0-1.0），默认使用配置中的 threshold

        Returns:
            是否溢出
        """
        if not self.config.max_tokens or self.config.max_tokens <= 0:
            return False

        tokens = self.get_token_count()
        if threshold is None:
            threshold = self.config.threshold
        max_threshold = self.config.max_tokens * threshold
        return tokens > max_threshold

    def _estimate_tokens(self) -> int:
        """估算 token 数量（字符数 / 4）"""
        return self._estimate_message_tokens(self._messages)

    async def _compress_async(self) -> None:
        """
        异步压缩上下文

        四级独立压缩策略：
        1. 弱 (Prune): 清理旧工具输出
        2. 中 (Summarize-Weak): 保留 3 轮对话
        3. 强 (Summarize-Strong): 只保留 1 轮对话
        4. 最强 (Summarize-All): 将会话全部压缩为一条消息

        每一级压缩都基于原始上下文生成新的消息列表，
        如果新消息列表不溢出则替换，否则尝试下一级。
        """
        try:
            if not self.config.max_tokens or self.config.max_tokens <= 0:
                log_error("context_compress_skip: invalid max_tokens")
                return

            # 保存原始消息，用于各级独立压缩
            original_messages = list(self._messages)

            log_event(
                "context_compress_start",
                total_messages=len(original_messages),
                original_token_count=self.get_token_count(),
            )

            # 阶段 1: 弱压缩 (Prune)
            pruned_messages = self._prune_tool_outputs(original_messages)
            token_count = await self._count_tokens_for_messages(pruned_messages)

            if token_count <= self.config.max_tokens * self.config.prune_threshold:
                self._messages = pruned_messages
                self._token_count = token_count
                log_event(
                    "context_compress_success",
                    strategy="prune",
                    final_token_count=token_count,
                    final_messages=len(pruned_messages),
                )
                await self._finish_compress()
                return

            log_event(
                "context_compress_upgrade",
                from_strategy="prune",
                to_strategy="summarize_weak",
                reason=f"prune_threshold={self.config.prune_threshold}",
            )

            # 阶段 2: 中压缩 (Summarize-Weak)
            summarized_weak_messages = await self._summarize_messages(
                original_messages,
                turns=self.config.summarize_weak_turns,
                min_msgs=self.config.summarize_weak_min_msgs,
                strategy="summarize_weak",
            )
            token_count = await self._count_tokens_for_messages(summarized_weak_messages)

            if token_count <= self.config.max_tokens * self.config.summarize_weak_threshold:
                self._messages = summarized_weak_messages
                self._token_count = token_count
                log_event(
                    "context_compress_success",
                    strategy="summarize_weak",
                    final_token_count=token_count,
                    final_messages=len(summarized_weak_messages),
                )
                await self._finish_compress()
                return

            log_event(
                "context_compress_upgrade",
                from_strategy="summarize_weak",
                to_strategy="summarize_strong",
                reason=f"summarize_weak_threshold={self.config.summarize_weak_threshold}",
            )

            # 阶段 3: 强压缩 (Summarize-Strong)
            summarized_strong_messages = await self._summarize_messages(
                original_messages,
                turns=self.config.summarize_strong_turns,
                min_msgs=self.config.summarize_strong_min_msgs,
                strategy="summarize_strong",
            )
            token_count = await self._count_tokens_for_messages(summarized_strong_messages)

            if token_count <= self.config.max_tokens * self.config.summarize_strong_threshold:
                self._messages = summarized_strong_messages
                self._token_count = token_count
                log_event(
                    "context_compress_success",
                    strategy="summarize_strong",
                    final_token_count=token_count,
                    final_messages=len(summarized_strong_messages),
                )
                await self._finish_compress()
                return

            log_event(
                "context_compress_upgrade",
                from_strategy="summarize_strong",
                to_strategy="summarize_all",
                reason=f"summarize_strong_threshold={self.config.summarize_strong_threshold}",
            )

            # 阶段 4: 最强压缩 (Summarize-All) - 兜底策略
            summarized_all_messages = await self._summarize_all(original_messages)
            token_count = await self._count_tokens_for_messages(summarized_all_messages)

            self._messages = summarized_all_messages
            self._token_count = token_count

            log_event(
                "context_compress_success",
                strategy="summarize_all",
                final_token_count=token_count,
                final_messages=len(summarized_all_messages),
            )

            await self._finish_compress()

        except Exception as e:
            log_error(f"context_compress_error: {str(e)}")
        finally:
            self._compressing = False

    async def _finish_compress(self) -> None:
        """完成压缩后的收尾工作"""
        # 压缩方法已经计算了 token 估算值，直接使用
        # 只在需要精确值时重新计算（可选）
        
        # 压缩后重写磁盘文件
        if self._context_file:
            await self._rewrite_to_disk()

    async def _count_tokens_for_messages(self, messages: List[Message]) -> int:
        """计算消息 token 数，优先使用 SDK 精确计数"""
        if self._count_tokens_fn:
            try:
                count = await self._count_tokens_fn(messages, self._system_prompts)
                if count > 0:
                    return count
            except Exception as e:
                log_error(f"context_count_tokens_error: {e}")

        return self._estimate_message_tokens(messages)

    def _prune_tool_outputs(self, messages: List[Message]) -> List[Message]:
        """
        Prune 阶段：清理旧的工具输出

        策略：
        1. 保护最近 2 轮对话的工具输出
        2. 保护最近 prune_protect_chars 字符的工具输出
        3. 超过阈值的旧工具输出替换为占位符

        Args:
            messages: 待压缩的消息列表

        Returns:
            压缩后的消息列表
        """
        pruned_chars = 0
        turn_count = 0

        # 从后往前遍历，统计工具输出信息
        tool_outputs_info = []  # [(msg_index, turn_count, char_count)]

        for i, msg in enumerate(reversed(messages)):
            orig_idx = len(messages) - 1 - i
            if msg.role == MessageRole.USER:
                turn_count += 1
            if msg.tool_result and msg.tool_result.content:
                char_count = len(msg.tool_result.content)
                tool_outputs_info.append((orig_idx, turn_count, char_count))

        # 确定哪些工具输出需要清理
        to_prune_indices = set()
        accumulated = 0

        for orig_idx, tc, char_count in tool_outputs_info:
            # 保护最近 N 条消息
            if orig_idx >= len(messages) - self.config.prune_keep_recent:
                accumulated += char_count
                continue

            # 保护最近 2 轮对话
            if tc < 2:
                accumulated += char_count
                continue
            # 保护最近 prune_protect_chars 字符
            if accumulated < self.config.prune_protect_chars:
                accumulated += char_count
                continue
            # 标记为需要清理
            to_prune_indices.add(orig_idx)
            pruned_chars += char_count

        # 只有当清理量超过最小阈值时才执行清理
        if pruned_chars < self.config.prune_min_chars:
            return list(messages)

        # 替换工具输出为占位符
        new_messages = []
        for i, msg in enumerate(messages):
            if i in to_prune_indices and msg.tool_result:
                new_messages.append(Message(
                    role=msg.role,
                    content=msg.content,
                    tool_calls=msg.tool_calls,
                    tool_result=ToolResult(
                        tool_call_id=msg.tool_result.tool_call_id,
                        name=msg.tool_result.name,
                        content=self.TOOL_OUTPUT_PLACEHOLDER,
                        is_error=msg.tool_result.is_error,
                    )
                ))
            else:
                new_messages.append(msg)

        log_event(
            "context_pruned",
            pruned_count=len(to_prune_indices),
            pruned_chars=pruned_chars,
        )

        return new_messages

    async def _summarize_messages(
        self,
        messages: List[Message],
        turns: int,
        min_msgs: int,
        strategy: str,
    ) -> List[Message]:
        """
        Summarize 阶段：使用 LLM 生成摘要

        Args:
            messages: 待压缩的消息列表
            turns: 保留最近 N 轮对话
            min_msgs: 最小保留消息数
            strategy: 策略名称 (summarize_weak / summarize_strong)

        Returns:
            压缩后的消息列表
        """
        if len(messages) <= min_msgs + 1:
            return list(messages)

        log_event(
            "context_summarize_start",
            total_messages=len(messages),
            strategy=strategy,
            keep_turns=turns,
            keep_min_msgs=min_msgs,
        )

        # 步骤 1: 按轮次分析消息
        turns_list = self._analyze_turns(messages)
        total_turns = len(turns_list)

        log_event(
            "context_summarize_turns_analysis",
            total_turns=total_turns,
            turns_detail=[len(t) for t in turns_list],
        )

        # 步骤 2: 保留最近 N 轮对话
        keep_turn_count = min(turns, total_turns)
        kept_msgs = []
        old_msgs = []

        if keep_turn_count > 0:
            # 保留最近的 N 轮
            for turn in turns_list[-keep_turn_count:]:
                kept_msgs.extend(turn)
            # 剩余的是旧消息
            for turn in turns_list[:-keep_turn_count]:
                old_msgs.extend(turn)
        else:
            # 没有足够的轮次，全部作为旧消息
            old_msgs = list(messages)
            kept_msgs = []

        # 步骤 3: 如果保留消息数 < min_msgs，扩展保留范围
        expanded_from_old = False
        if len(kept_msgs) < min_msgs:
            # 从旧消息中补充
            needed = min_msgs - len(kept_msgs)
            if old_msgs:
                # 取旧消息的末尾部分
                additional = old_msgs[-needed:] if len(old_msgs) > needed else old_msgs
                kept_msgs = additional + kept_msgs
                old_msgs = old_msgs[:-len(additional)] if len(old_msgs) > len(additional) else []
                expanded_from_old = len(additional) > 0

        log_event(
            "context_summarize_after_turns",
            kept_count=len(kept_msgs),
            old_count=len(old_msgs),
        )

        # 步骤 4: 移除不完整的 tool_calls
        # 只有当保留消息数 < min_msgs（需要从旧消息补充）时，才可能有不完整的 tool_calls
        if expanded_from_old:
            kept_msgs = self._remove_incomplete_tool_calls(kept_msgs)

        # 步骤 5: 如果没有旧消息需要摘要，直接完成
        if not old_msgs:
            log_event(
                "context_summarize_no_old_messages",
                strategy=strategy,
                kept_count=len(kept_msgs),
            )
            return kept_msgs

        # 步骤 6: 使用 LLM 生成摘要
        summary = await self._generate_summary(old_msgs)

        # 步骤 7: 重建消息列表
        new_messages = []

        # 将摘要作为用户消息添加
        new_messages.append(Message(
            role=MessageRole.USER,
            content=f"<context-summary>\n{summary}\n</context-summary>"
        ))

        # 添加助手确认消息，保持对话流畅
        new_messages.append(Message(
            role=MessageRole.ASSISTANT,
            content="好的，我已了解之前的对话内容，请继续。"
        ))

        # 添加保留的最近消息
        new_messages.extend(kept_msgs)

        log_event(
            "context_summarized",
            strategy=strategy,
            compressed_count=len(old_msgs),
            kept_count=len(new_messages),
            summary_length=len(summary),
        )

        return new_messages

    async def _summarize_all(self, messages: List[Message]) -> List[Message]:
        """
        最强压缩：将所有消息压缩为一条摘要消息

        适用于：即使保留 1 轮对话后，上下文仍然溢出

        Args:
            messages: 待压缩的消息列表

        Returns:
            压缩后的消息列表
        """
        if len(messages) <= 2:
            return list(messages)

        log_event(
            "context_summarize_all_start",
            total_messages=len(messages),
        )

        # 使用 LLM 生成摘要
        summary = await self._generate_summary(messages)

        # 重建消息列表：只有摘要 + 确认
        new_messages = []

        # 将摘要作为用户消息添加
        new_messages.append(Message(
            role=MessageRole.USER,
            content=f"<context-summary>\n{summary}\n</context-summary>"
        ))

        # 添加助手确认消息，保持对话流畅
        new_messages.append(Message(
            role=MessageRole.ASSISTANT,
            content="好的，我已了解之前的全部对话内容，请继续。"
        ))

        log_event(
            "context_summarize_all_completed",
            summary_length=len(summary),
        )

        return new_messages

    def _analyze_turns(self, messages: List[Message]) -> List[List[Message]]:
        """
        分析消息，按轮次分组

        轮次定义：从 User 消息开始，到下一个 User 消息之前结束
        每轮可能包含：User + Assistant[ + Tool Calls + Tool Results]

        Args:
            messages: 待分析的消息列表

        Returns:
            轮次列表，每个元素是一轮中的消息列表
        """
        turns: List[List[Message]] = []
        current_turn: List[Message] = []

        for msg in messages:
            # User 消息表示新一轮开始，先保存上一轮
            if msg.role == MessageRole.USER:
                if current_turn:
                    turns.append(current_turn)
                    current_turn = []

            current_turn.append(msg)

        # 保存最后一轮
        if current_turn:
            turns.append(current_turn)

        return turns

    def _remove_incomplete_tool_calls(self, messages: List[Message]) -> List[Message]:
        """
        移除不完整的 tool_calls（无对应 tool_result）

        策略：如果一条 Assistant 消息有 tool_calls 但没有对应的 tool_result，
        清空 tool_calls 字段以避免 LLM 困惑

        Returns:
            处理后的消息列表
        """
        # 建立 tool_call_id -> tool_result 的映射
        tool_result_ids = set()
        for msg in messages:
            if msg.tool_result and msg.tool_result.tool_call_id:
                tool_result_ids.add(msg.tool_result.tool_call_id)

        # 清理不完整的 tool_calls
        cleaned_messages = []
        removed_count = 0

        for msg in messages:
            if msg.tool_calls:
                # 只保留有对应 tool_result 的调用，移除不完整调用
                matched_tool_calls = [
                    tc for tc in msg.tool_calls
                    if tc.id
                    and tc.id in tool_result_ids
                ]
                if len(matched_tool_calls) != len(msg.tool_calls):
                    cleaned_msg = Message(
                        role=msg.role,
                        content=msg.content,
                        tool_calls=matched_tool_calls,
                        tool_result=msg.tool_result,
                    )
                    cleaned_messages.append(cleaned_msg)
                    removed_count += 1
                    continue
            cleaned_messages.append(msg)

        if removed_count > 0:
            log_event(
                "context_incomplete_tool_calls_removed",
                count=removed_count,
            )

        return cleaned_messages

    def _estimate_message_tokens(self, messages: List[Message]) -> int:
        """估算消息列表的 token 数（使用字符估算）"""
        total_chars = sum(len(m.content) for m in self._system_prompts if m.content)
        for msg in messages:
            if msg.content:
                total_chars += len(msg.content)
            for tc in (msg.tool_calls or []):
                total_chars += len(str(tc.arguments))
            if msg.tool_result:
                total_chars += len(msg.tool_result.content or "")
        return total_chars // self.config.chars_per_token

    async def _generate_summary(self, messages: List[Message]) -> str:
        """使用 LLM 生成消息摘要"""
        formatted = self._format_messages_for_summary(messages)

        summary_request = f"""请为继续对话提供一个详细的摘要。重点关注：
1. 我们做了什么
2. 我们正在做什么
3. 我们正在处理哪些文件
4. 接下来我们要做什么

对话历史：
{formatted}

请用中文生成摘要，控制在 500 字以内："""

        try:
            summary = ""
            async for chunk in self.chat_fn(
                messages=[
                    Message(role=MessageRole.USER, content=summary_request),
                ],
                system_prompts=[
                    Message(role=MessageRole.SYSTEM, content=self._compaction_prompt),
                ],
                tools=None,
            ):
                if chunk.text:
                    summary += chunk.text

            if summary:
                return summary
            else:
                return self._fallback_summary(messages)

        except Exception as e:
            log_error(f"context_summary_error: {str(e)}")
            return self._fallback_summary(messages)

    def _format_messages_for_summary(self, messages: List[Message]) -> str:
        """格式化消息为文本用于摘要"""
        lines = []
        for msg in messages:
            role_name = {
                MessageRole.USER: "用户",
                MessageRole.ASSISTANT: "助手",
                MessageRole.TOOL: "工具",
            }.get(msg.role, str(msg.role))

            if msg.tool_calls:
                if msg.content:
                    content = msg.content[:300]
                    if len(msg.content) > 300:
                        content += "..."
                    lines.append(f"[{role_name}] {content}")
                for tc in msg.tool_calls:
                    lines.append(f"[{role_name}] 调用工具: {tc.name}")
            elif msg.tool_result:
                status = "错误" if msg.tool_result.is_error else "成功"
                lines.append(f"[工具 {msg.tool_result.name}] 执行{status}")
            elif msg.content:
                content = msg.content[:500]
                if len(msg.content) > 500:
                    content += "..."
                lines.append(f"[{role_name}] {content}")

        return "\n".join(lines)

    def _fallback_summary(self, messages: List[Message]) -> str:
        """降级摘要（当 LLM 调用失败时使用）"""
        user_messages = [m for m in messages if m.role == MessageRole.USER and m.content]
        tool_calls = []
        for m in messages:
            if m.tool_calls:
                tool_calls.extend([tc.name for tc in m.tool_calls])

        parts = []
        parts.append(f"共 {len(messages)} 条历史消息")

        if user_messages:
            recent = user_messages[-3:] if len(user_messages) > 3 else user_messages
            user_summary = "; ".join([m.content[:100] for m in recent])
            parts.append(f"用户请求: {user_summary}")

        if tool_calls:
            unique_tools = list(set(tool_calls))
            parts.append(f"使用的工具: {', '.join(unique_tools)}")

        return "\n".join(parts)

    # ==================== 持久化方法 ====================

    async def _append_to_disk(self, message: Message) -> None:
        """追加单条消息到磁盘（JSONL 格式）"""
        async with self._write_lock:
            try:
                line = json.dumps(message.to_dict(), ensure_ascii=False) + "\n"
                await asyncio.to_thread(self._sync_append, line)
            except Exception as e:
                log_error(f"context_persist_error: {str(e)}")

    def _sync_append(self, line: str) -> None:
        """同步追加写入"""
        with open(self._context_file, "a", encoding="utf-8") as f:
            f.write(line)

    async def _rewrite_to_disk(self) -> None:
        """重写整个文件（压缩后调用），先写 tmp 再 rename 保证原子性"""
        async with self._write_lock:
            try:
                lines = [
                    json.dumps(msg.to_dict(), ensure_ascii=False) + "\n"
                    for msg in self._messages
                ]
                tmp_file = self._context_file.with_suffix(".jsonl.tmp")
                await asyncio.to_thread(self._sync_rewrite, tmp_file, lines)
                tmp_file.rename(self._context_file)
            except Exception as e:
                log_error(f"context_rewrite_error: {str(e)}")

    def _sync_rewrite(self, path: Path, lines: List[str]) -> None:
        """同步重写文件"""
        with open(path, "w", encoding="utf-8") as f:
            f.writelines(lines)

    async def load_from_disk(self) -> bool:
        """从磁盘恢复消息列表"""
        if not self._context_file or not self._context_file.exists():
            return False

        try:
            messages = await asyncio.to_thread(self._sync_load)
            self._messages = messages
            self._token_count = self._estimate_tokens()
            log_event("context_loaded_from_disk", message_count=len(messages))
            return True
        except Exception as e:
            log_error(f"context_load_error: {str(e)}")
            return False

    def _sync_load(self) -> List[Message]:
        """同步读取 JSONL 文件"""
        messages = []
        with open(self._context_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    data = json.loads(line)
                    messages.append(Message.from_dict(data))
        return messages
