"""
主循环级上下文管理（Prune & Compress，占位）。

当前阶段只实现 **Prune**，用于在每轮模型调用前，对历史 ToolMessage
进行“轻量瘦身”，避免特别久远的大量工具输出持续占用上下文。

设计要点（与 `.claude/context/agent_loop_prune_and_compress.md` 对齐）：
- 只动 ToolMessage 的 content + metadata，不删除消息本身；
- 保护最近两轮用户对话及其工具调用；
- 保护最近约 PRUNE_PROTECT_TOKENS 的工具输出；
- 保护关键工具（如 `skill`）的输出；
- 只有当可剪枝总量 > PRUNE_MINIMUM_TOKENS 时才真正执行 compact。

注意：
- 本模块不依赖 LangGraph 具体结构，只接收/返回 state 片段；
- 压缩（Compress / summary）逻辑留待后续阶段扩展。
"""

from __future__ import annotations

import os
import time
from typing import Any, Dict, List, Tuple

from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage
from langchain_core.messages.utils import count_tokens_approximately

# 粗略估算：约 4 个字符 ≈ 1 token（中英文混合场景下的经验值）
NUM_CHARS_PER_TOKEN = 4

# 最近这段工具输出（按 token 计）始终保留，不参与剪枝
PRUNE_PROTECT_TOKENS = 40_000

# 只有“候选被剪枝的旧工具输出”累计超过该阈值时，才真正执行 compact
PRUNE_MINIMUM_TOKENS = 20_000

# 长期重要的工具（例如 server_skill_middleware 注入的 skill 工具）
PRUNE_PROTECTED_TOOLS = frozenset({"skill"})

# Compress 相关默认参数（均可通过环境变量覆盖）
DEFAULT_MAX_CONTEXT_TOKENS = 120_000
DEFAULT_COMPRESS_TRIGGER_FRACTION = 0.8
DEFAULT_COMPRESS_KEEP_RECENT_MESSAGES = 8


def _estimate_tokens_from_text(content: str) -> int:
    """根据文本长度粗略估算 token 数量。"""
    if not content:
        return 0
    return max(1, len(content) // NUM_CHARS_PER_TOKEN)


def _is_summary_message(msg: BaseMessage) -> bool:
    """判断是否为 summary 消息（用于作为 prune 边界）。

    说明：
    - Compress 阶段生成的是 HumanMessage(summary=True)；
    - 未来如需改为 AIMessage(summary=True)，这里同样兼容。
    """
    if not isinstance(msg, (AIMessage, HumanMessage)):
        return False
    extra: Dict[str, Any] = getattr(msg, "additional_kwargs", {}) or {}
    metadata: Dict[str, Any] = getattr(msg, "metadata", {}) or {}
    return bool(extra.get("summary") or metadata.get("summary"))


def _get_tool_name(msg: ToolMessage, messages: List[BaseMessage] | None = None) -> str | None:
    """尽量从 ToolMessage 中解析出工具名称。

    兜底逻辑：当 name 缺失时（如 wrap_tool_call 缓存返回的 ToolMessage），
    从 messages 中向前查找 AIMessage.tool_calls 里匹配 tool_call_id 的 name。
    """
    extra: Dict[str, Any] = getattr(msg, "additional_kwargs", {}) or {}
    metadata: Dict[str, Any] = getattr(msg, "metadata", {}) or {}

    name = (
        getattr(msg, "tool_name", None)
        or getattr(msg, "name", None)
        or extra.get("tool_name")
        or extra.get("name")
        or metadata.get("tool_name")
        or metadata.get("name")
    )
    if name is not None:
        return name

    # 兜底：从 messages 中通过 tool_call_id 反查（应对 wrap_tool_call 缓存等未传 name 的场景）
    if messages and (tid := getattr(msg, "tool_call_id", None)):
        try:
            idx = messages.index(msg)
        except ValueError:
            idx = -1
        for i in range(idx - 1, -1, -1):
            m = messages[i]
            if isinstance(m, AIMessage) and m.tool_calls:
                for tc in m.tool_calls:
                    if isinstance(tc, dict) and tc.get("id") == tid:
                        return tc.get("name")
    return None


def _is_compacted(msg: ToolMessage) -> bool:
    """是否已经被 compact 过。"""
    extra: Dict[str, Any] = getattr(msg, "additional_kwargs", {}) or {}
    metadata: Dict[str, Any] = getattr(msg, "metadata", {}) or {}
    return bool(extra.get("compacted") or metadata.get("compacted"))


def _mark_compacted(msg: ToolMessage, original_length: int) -> None:
    """在 ToolMessage 上打 compact 标记并替换内容为占位文本（就地修改）。"""
    placeholder = "[Old tool result compacted]"

    msg.content = placeholder

    # 追加 compact 标记到 additional_kwargs / metadata
    extra: Dict[str, Any] = getattr(msg, "additional_kwargs", None)
    if extra is None:
        extra = {}
        setattr(msg, "additional_kwargs", extra)

    metadata: Dict[str, Any] = getattr(msg, "metadata", None)
    if metadata is None:
        metadata = {}
        setattr(msg, "metadata", metadata)

    now_ms = int(time.time() * 1000)

    extra.setdefault("compacted", True)
    extra.setdefault("compacted_at", now_ms)
    extra.setdefault("original_length", original_length)

    metadata.setdefault("compacted", True)
    metadata.setdefault("compacted_at", now_ms)
    metadata.setdefault("original_length", original_length)


def _get_int_env(name: str, default: int) -> int:
    """读取整型环境变量，非法时回退到默认值。"""
    raw = os.getenv(name)
    if raw is None:
        return default
    try:
        value = int(raw)
        if value > 0:
            return value
    except Exception:
        return default
    return default


def _get_float_env(name: str, default: float) -> float:
    """读取浮点环境变量，非法时回退到默认值。"""
    raw = os.getenv(name)
    if raw is None:
        return default
    try:
        value = float(raw)
        if value > 0:
            return value
    except Exception:
        return default
    return default


def _collect_prune_candidates(messages: List[BaseMessage]) -> Tuple[List[int], int]:
    """从消息列表中收集需要 compact 的 ToolMessage 索引及累计 token 数。

    返回：
        (to_compact_indices, pruned_tokens)
    """
    to_compact_indices: List[int] = []
    pruned_tokens = 0
    total_tokens_recent_tools = 0
    turns = 0  # 用户轮次（HumanMessage 数量）

    # 从最新消息开始向前扫描
    for idx in range(len(messages) - 1, -1, -1):
        msg = messages[idx]

        # 1. 保护最近两轮用户对话
        if isinstance(msg, HumanMessage):
            turns += 1
            if turns < 2:
                # 最近两轮 user + 其工具全保护，不剪
                continue

        # 2. summary 边界
        if _is_summary_message(msg):
            break

        # 3. 只处理 ToolMessage
        if not isinstance(msg, ToolMessage):
            continue

        # 3.1 保护关键工具（例如 skill）
        tool_name = _get_tool_name(msg, messages)
        if tool_name in PRUNE_PROTECTED_TOOLS:
            continue

        # 3.2 已经被 compact 过：视为边界，停止向前扫描
        if _is_compacted(msg):
            break

        # 3.3 统计 token 数
        content = msg.content or ""
        if not isinstance(content, str):
            continue

        estimate = _estimate_tokens_from_text(content)
        total_tokens_recent_tools += estimate

        # 3.4 保护“最近约 PRUNE_PROTECT_TOKENS 的工具输出”
        if total_tokens_recent_tools <= PRUNE_PROTECT_TOKENS:
            continue

        # 3.5 其余更早的工具输出进入剪枝候选集合
        pruned_tokens += estimate
        to_compact_indices.append(idx)

    return to_compact_indices, pruned_tokens


DEFAULT_COMPACTION_PROMPT = """Provide a detailed prompt for continuing our conversation above.
Focus on information that would be helpful for continuing the conversation, including what we did, what we're doing, which files we're working on, and what we're going to do next.
The summary that you construct will be used so that another agent can read it and continue the work.

When constructing the summary, try to stick to this template:
---
## Goal

[What goal(s) is the user trying to accomplish?]

## Instructions

- [What important instructions did the user give you that are relevant]
- [If there is a plan or spec, include information about it so next agent can continue using it]

## Discoveries

[What notable things were learned during this conversation that would be useful for the next agent to know when continuing the work]

## Accomplished

[What work has been completed, what work is still in progress, and what work is left?]

## Relevant files / directories

[Construct a structured list of relevant files that have been read, edited, or created that pertain to the task at hand. If all the files in a directory are relevant, include the path to the directory.]
---"""


def _compress_messages_if_needed(
    messages: List[BaseMessage],
    model: Any | None = None,
) -> Tuple[List[BaseMessage], bool]:
    """在满足上下文阈值时，对历史消息做一次摘要压缩。

    当前实现为“结构化机械摘要”：
    - 使用 count_tokens_approximately 估算当前上下文体积；
    - 当超过触发阈值时，按 keep_recent_messages 保留末尾若干条；
    - 保留首条 HumanMessage（H0）及 skill 工具输出（若位于被压缩区间）；
    - 对被压缩区间构造一条带有 summary=True 标记的 HumanMessage 作为概要；
    - 返回新的消息列表。

    说明：
    - 若提供了 model，则优先使用当前 ChatModel 按 OpenCode 的 compaction prompt 生成摘要；
    - 若 model 为空或调用失败，则回退到“结构化机械摘要”。
    """
    if not messages:
        return messages, False

    # 读取 Compress 相关配置（均可通过环境变量覆盖）
    max_context_tokens = _get_int_env("OPENCODE_MAX_CONTEXT_TOKENS", DEFAULT_MAX_CONTEXT_TOKENS)
    trigger_fraction = _get_float_env(
        "OPENCODE_COMPRESS_TRIGGER_FRACTION",
        DEFAULT_COMPRESS_TRIGGER_FRACTION,
    )
    keep_recent_messages = _get_int_env(
        "OPENCODE_COMPRESS_KEEP_RECENT_MESSAGES",
        DEFAULT_COMPRESS_KEEP_RECENT_MESSAGES,
    )

    if trigger_fraction <= 0 or trigger_fraction >= 1:
        # 非法配置则不启用 Compress，避免意外行为
        return messages, False

    total_tokens = count_tokens_approximately(messages)
    trigger_tokens = int(max_context_tokens * trigger_fraction)

    # 未达到触发阈值，不做压缩
    if total_tokens < trigger_tokens:
        return messages, False

    n = len(messages)
    if n <= keep_recent_messages + 2:
        # 消息太少，压缩意义不大
        return messages, False

    cutoff_index = max(0, n - keep_recent_messages)

    # 收集被压缩区间（不含尾部 keep_recent_messages）的消息
    prefix = messages[:cutoff_index]
    suffix = messages[cutoff_index:]

    # 1) 找到首条 HumanMessage（H0）
    first_human: HumanMessage | None = None
    for msg in prefix:
        if isinstance(msg, HumanMessage):
            first_human = msg
            break

    # 2) 收集 skill 工具输出（仅在 prefix 中）
    skill_tool_messages: List[ToolMessage] = []
    for msg in prefix:
        if isinstance(msg, ToolMessage):
            tool_name = _get_tool_name(msg, messages)
            if tool_name in PRUNE_PROTECTED_TOOLS:
                skill_tool_messages.append(msg)

    # 3) 找到最新一次 plan 工具输出（例如 write_todos），仅在 prefix 中
    PLAN_TOOL_NAMES = frozenset({"write_todos", "plan"})
    last_plan_tool: ToolMessage | None = None
    last_plan_ai: AIMessage | None = None
    for idx, msg in enumerate(prefix):
        if isinstance(msg, ToolMessage):
            tool_name = _get_tool_name(msg, messages)
            if tool_name in PLAN_TOOL_NAMES:
                last_plan_tool = msg
                # 尝试关联其前面的 AIMessage 作为发起调用的消息
                if idx - 1 >= 0 and isinstance(prefix[idx - 1], AIMessage):
                    last_plan_ai = prefix[idx - 1]  # type: ignore[assignment]

    # 4) 构造摘要文本，优先使用 LLM；失败时回退到机械摘要
    # Goal: 使用 H0 的内容
    goal_text = (first_human.content if isinstance(first_human, HumanMessage) else "") or ""
    if isinstance(goal_text, str):
        goal_text = goal_text[:1000]
    else:
        goal_text = str(goal_text)[:1000]

    # Instructions: 其他 HumanMessage 的内容
    instructions_lines: List[str] = []
    for msg in prefix:
        if isinstance(msg, HumanMessage) and msg is not first_human:
            content = msg.content
            if not isinstance(content, str):
                content = str(content)
            instructions_lines.append(f"- {content[:500]}")

    # Discoveries: 取若干 AIMessage 的内容
    discoveries_lines: List[str] = []
    for msg in prefix:
        if isinstance(msg, AIMessage) and msg is not last_plan_ai:
            content = msg.content
            if not isinstance(content, str):
                content = str(content)
            if content:
                discoveries_lines.append(f"- {content[:500]}")

    # Accomplished: 简要描述已执行的工具调用数量
    compressed_tool_calls = [
        msg for msg in prefix if isinstance(msg, ToolMessage) and msg not in skill_tool_messages
    ]
    accomplished_lines: List[str] = []
    if compressed_tool_calls:
        accomplished_lines.append(
            f"- Executed {len(compressed_tool_calls)} tool calls in the compressed history."
        )

    # Relevant files / directories: 从 ToolMessage 文本中简单抽取 path 片段（非常粗糙）
    relevant_paths: List[str] = []
    for msg in prefix:
        if isinstance(msg, ToolMessage):
            content = msg.content
            if not isinstance(content, str):
                content = str(content)
            # 简单启发式：提取看起来像路径的子串
            for token in content.split():
                if "/" in token and len(token) > 2:
                    relevant_paths.append(token.strip('",\\\''))
    # 去重，保留顺序
    seen: set[str] = set()
    unique_paths: List[str] = []
    for p in relevant_paths:
        if p not in seen:
            seen.add(p)
            unique_paths.append(p)
    relevant_lines = [f"- {p}" for p in unique_paths[:20]]

    summary_text: str

    # 4.1 尝试使用当前模型生成摘要（路线 B）
    if model is not None:
        try:
            # 仅基于被压缩区间 prefix 生成摘要，避免重复考虑最近消息
            llm_messages: List[BaseMessage] = list(prefix)
            llm_messages.append(HumanMessage(content=DEFAULT_COMPACTION_PROMPT))
            llm_result = model.invoke(llm_messages)  # type: ignore[call-arg]
            content = getattr(llm_result, "content", None)
            if isinstance(content, str) and content.strip():
                summary_text = content
            else:
                # 若模型没有返回可用文本，则退回机械摘要
                raise ValueError("Empty summary content from model")
        except Exception:
            # 出现任何异常时，退回到机械摘要，以保证主循环稳定性
            summary_text = ""
    else:
        summary_text = ""

    # 4.2 机械摘要（作为 fallback，也用于 model 不可用的场景）
    if not summary_text:
        summary_text_parts: List[str] = []
        summary_text_parts.append("## Goal\n\n" + (goal_text or "[unknown]"))
        summary_text_parts.append(
            "## Instructions\n\n"
            + ("\n".join(instructions_lines) if instructions_lines else "[none]")
        )
        summary_text_parts.append(
            "## Discoveries\n\n" + ("\n".join(discoveries_lines) if discoveries_lines else "[none]")
        )
        summary_text_parts.append(
            "## Accomplished\n\n"
            + ("\n".join(accomplished_lines) if accomplished_lines else "[none]")
        )
        summary_text_parts.append(
            "## Relevant files / directories\n\n"
            + ("\n".join(relevant_lines) if relevant_lines else "[none]")
        )
        summary_text = "\n\n".join(summary_text_parts)

    summary_msg = HumanMessage(
        content=summary_text,
        additional_kwargs={"summary": True},
    )

    # 5) 重建消息列表：保留头部关键消息 + 摘要 + 尾部最近若干消息
    new_messages: List[BaseMessage] = []

    # 5.1 首条 HumanMessage（H0）
    if first_human is not None:
        new_messages.append(first_human)

    # 5.2 skill 工具输出（仅在 prefix 中且尚未添加）
    for msg in skill_tool_messages:
        if msg not in new_messages:
            new_messages.append(msg)

    # 5.3 最新 plan AI + Tool（仅在 prefix 中且尚未添加）
    if last_plan_ai is not None and last_plan_ai not in new_messages:
        new_messages.append(last_plan_ai)
    if last_plan_tool is not None and last_plan_tool not in new_messages:
        new_messages.append(last_plan_tool)

    # 5.4 摘要消息
    new_messages.append(summary_msg)

    # 5.5 追加尾部最近 keep_recent_messages 条原始消息
    new_messages.extend(suffix)

    return new_messages, True


def prune_and_compress_state(state: Dict[str, Any], model: Any | None = None) -> Dict[str, Any]:
    """对 state 执行 Compress + Prune。

    执行顺序：
    1. Compress：在上下文接近上限时，对较早历史插入一条摘要消息；
    2. Prune：在（可能已被 Compress 过的）消息列表上，对旧 ToolMessage 进行占位裁剪。

    返回值是“增量更新”字典（通常只包含更新后的 messages），
    以便作为 LangGraph 节点的 `Command(update=...)` 使用。
    """
    messages = state.get("messages")
    if not isinstance(messages, list) or not messages:
        return {}

    # 1) Compress：在上下文超阈值时插入摘要消息
    working_messages: List[BaseMessage] = list(messages)
    working_messages, did_compress = _compress_messages_if_needed(working_messages, model=model)

    # 2) Prune：在（可能包含 summary）的消息列表上执行裁剪
    to_compact_indices, pruned_tokens = _collect_prune_candidates(working_messages)

    # 没有候选或收益太小 → 只在发生了压缩时返回 Compress 结果
    if not to_compact_indices or pruned_tokens <= PRUNE_MINIMUM_TOKENS:
        if did_compress and working_messages != messages:
            return {"messages": working_messages}
        return {}

    # 复制列表，避免直接修改原列表引用（消息对象本身就地修改）
    new_messages: List[BaseMessage] = list(working_messages)

    for idx in to_compact_indices:
        msg = new_messages[idx]
        if not isinstance(msg, ToolMessage):
            continue
        content = msg.content or ""
        original_length = len(content) if isinstance(content, str) else 0
        _mark_compacted(msg, original_length)

    return {"messages": new_messages}


__all__ = [
    "prune_and_compress_state",
    "PRUNE_PROTECT_TOKENS",
    "PRUNE_MINIMUM_TOKENS",
    "PRUNE_PROTECTED_TOOLS",
]

