cmdc

defmodule CMDC.Agent.Compactor.ArgTruncator do
  @moduledoc """
  Tool 参数预截断 — 对标 DeepAgents `TruncateArgsSettings`。

  这是一个 **轻量级的 pre-summarization 优化**，在完整压缩之前运行。
  当 token 使用量超过配置阈值时，自动截断保留窗口之外的旧消息中
  `write_file`、`edit_file`、`execute` 等工具调用的大参数字段。

  ## 工作原理

  1. 检测触发条件（token 数 / 消息数 / fraction）
  2. 计算保留窗口——最近的 N 条消息不动
  3. 在保留窗口之前的 assistant 消息中，扫描 tool_calls
  4. 对匹配工具名的 tool_call arguments 中超长字符串值进行截断

  ## 截断规则

  - 只处理 `arguments` map 中值为字符串且超过 `max_length` 的字段
  - 截断为：前 20 字符 + `truncation_text`（默认 `"...(argument truncated)"`）
  - 仅修改历史消息的 **副本**，不影响原始数据

  ## 配置

  通过 `State.config` 传入：

  - `:truncate_args_trigger` — `{:tokens, n}` | `{:messages, n}` | `{:fraction, f}` | `nil`（禁用）
  - `:truncate_args_keep` — `{:messages, n}` | `{:fraction, f}`（默认 `{:messages, 20}`）
  - `:truncate_args_max_length` — 单个参数最大字符数（默认 2000）
  - `:truncate_args_text` — 截断后缀文本
  - `:truncate_args_tools` — 要截断的工具名列表（默认 `["write_file", "edit_file", "execute"]`）
  """

  alias CMDC.Agent.State

  # ==========================================================================
  # 默认配置
  # ==========================================================================

  @default_keep {:messages, 20}
  @default_max_length 2_000
  @default_truncation_text "...(argument truncated)"
  @default_target_tools MapSet.new(["write_file", "edit_file", "execute"])
  @preview_length 20
  @chars_per_token 4

  # ==========================================================================
  # 类型定义
  # ==========================================================================

  @typedoc "截断触发配置。`nil` 表示禁用。"
  @type trigger ::
          {:tokens, pos_integer()} | {:messages, pos_integer()} | {:fraction, float()} | nil

  @typedoc "保留策略。"
  @type keep :: {:messages, pos_integer()} | {:fraction, float()}

  @typedoc "截断配置。"
  @type truncate_opts :: %{
          trigger: trigger(),
          keep: keep(),
          max_length: pos_integer(),
          truncation_text: String.t(),
          target_tools: MapSet.t(String.t())
        }

  # ==========================================================================
  # Public API
  # ==========================================================================

  @doc """
  对 State 中的消息执行参数截断。

  返回更新后的 State（消息可能被修改），以及是否有实际修改。

  - `{state, true}` — 至少一条消息被截断
  - `{state, false}` — 无需截断或未达到阈值
  """
  @spec truncate(State.t()) :: {State.t(), boolean()}
  def truncate(%State{} = state) do
    opts = truncate_config(state)

    if opts.trigger == nil do
      {state, false}
    else
      do_truncate(state, opts)
    end
  end

  @doc """
  读取截断配置。从 `state.config` 中获取，缺省使用默认值。
  """
  @spec truncate_config(State.t()) :: truncate_opts()
  def truncate_config(%State{config: config}) when is_map(config) do
    target_tools =
      case Map.get(config, :truncate_args_tools) do
        nil -> @default_target_tools
        tools when is_list(tools) -> MapSet.new(tools)
        %MapSet{} = set -> set
      end

    %{
      trigger: Map.get(config, :truncate_args_trigger),
      keep: Map.get(config, :truncate_args_keep, @default_keep),
      max_length: Map.get(config, :truncate_args_max_length, @default_max_length),
      truncation_text: Map.get(config, :truncate_args_text, @default_truncation_text),
      target_tools: target_tools
    }
  end

  def truncate_config(_state) do
    %{
      trigger: nil,
      keep: @default_keep,
      max_length: @default_max_length,
      truncation_text: @default_truncation_text,
      target_tools: @default_target_tools
    }
  end

  @doc """
  截断单个 tool_call 的大参数。

  返回 `{new_tool_call, modified?}`。
  """
  @spec truncate_tool_call(map(), pos_integer(), String.t()) :: {map(), boolean()}
  def truncate_tool_call(tool_call, max_length, truncation_text) do
    args = tool_call[:arguments] || tool_call["arguments"] || %{}

    {new_args, modified} =
      Enum.reduce(args, {%{}, false}, fn {key, value}, {acc, mod} ->
        if is_binary(value) and byte_size(value) > max_length do
          truncated = String.slice(value, 0, @preview_length) <> truncation_text
          {Map.put(acc, key, truncated), true}
        else
          {Map.put(acc, key, value), mod}
        end
      end)

    if modified do
      new_tc = put_arguments(tool_call, new_args)
      {new_tc, true}
    else
      {tool_call, false}
    end
  end

  # ==========================================================================
  # 私有 — 截断执行
  # ==========================================================================

  defp do_truncate(state, opts) do
    total_tokens = estimate_tokens(state)
    chronological = Enum.reverse(state.messages)

    if should_truncate?(chronological, total_tokens, opts, state) do
      apply_truncation(state, chronological, opts)
    else
      {state, false}
    end
  end

  defp apply_truncation(state, chronological, opts) do
    cutoff = determine_cutoff(chronological, opts, state)

    if cutoff >= length(chronological) do
      {state, false}
    else
      {new_msgs, modified} = truncate_messages(chronological, cutoff, opts)

      if modified do
        {%{state | messages: Enum.reverse(new_msgs)}, true}
      else
        {state, false}
      end
    end
  end

  # ==========================================================================
  # 私有 — 阈值判定
  # ==========================================================================

  defp should_truncate?(messages, total_tokens, opts, state) do
    case opts.trigger do
      {:messages, n} ->
        length(messages) >= n

      {:tokens, n} ->
        total_tokens >= n

      {:fraction, f} when is_float(f) ->
        context_window = state.token_usage.context_window

        if context_window > 0 do
          threshold = trunc(context_window * f)
          total_tokens >= max(threshold, 1)
        else
          false
        end

      _ ->
        false
    end
  end

  # ==========================================================================
  # 私有 — 保留窗口计算
  # ==========================================================================

  defp determine_cutoff(messages, opts, state) do
    total = length(messages)

    case opts.keep do
      {:messages, n} ->
        if total <= n, do: total, else: total - n

      {:fraction, f} when is_float(f) ->
        keep_count = count_keep_for_fraction(messages, f, state)
        if total <= keep_count, do: total, else: total - keep_count
    end
  end

  defp count_keep_for_fraction(messages, f, state) do
    context_window = state.token_usage.context_window

    target_tokens =
      if context_window > 0 do
        trunc(context_window * f)
      else
        20 * 250
      end

    count_keep_from_end(Enum.reverse(messages), target_tokens)
  end

  defp count_keep_from_end(reversed_msgs, target_tokens) do
    {count, _tokens} =
      Enum.reduce_while(reversed_msgs, {0, 0}, fn msg, {count, tokens} ->
        msg_tokens = div(byte_size(msg.content || ""), @chars_per_token)
        new_tokens = tokens + msg_tokens

        if new_tokens > target_tokens do
          {:halt, {count, new_tokens}}
        else
          {:cont, {count + 1, new_tokens}}
        end
      end)

    max(count, 1)
  end

  # ==========================================================================
  # 私有 — 消息遍历截断
  # ==========================================================================

  defp truncate_messages(chronological, cutoff, opts) do
    {processed, any_modified} =
      chronological
      |> Enum.with_index()
      |> Enum.map_reduce(false, fn {msg, idx}, modified ->
        truncate_message_at(msg, idx, cutoff, opts, modified)
      end)

    {processed, any_modified}
  end

  defp truncate_message_at(msg, idx, cutoff, opts, modified) do
    if idx < cutoff and msg.role == :assistant and is_list(msg.tool_calls) and
         msg.tool_calls != [] do
      {new_calls, msg_modified} = truncate_tool_calls(msg.tool_calls, opts)

      if msg_modified do
        {%{msg | tool_calls: new_calls}, true}
      else
        {msg, modified}
      end
    else
      {msg, modified}
    end
  end

  defp truncate_tool_calls(tool_calls, opts) do
    Enum.map_reduce(tool_calls, false, fn tc, modified ->
      truncate_single_tool_call(tc, opts, modified)
    end)
  end

  defp truncate_single_tool_call(tc, opts, modified) do
    name = tc[:name] || tc["name"] || ""

    if MapSet.member?(opts.target_tools, name) do
      {new_tc, tc_modified} = truncate_tool_call(tc, opts.max_length, opts.truncation_text)
      {new_tc, modified or tc_modified}
    else
      {tc, modified}
    end
  end

  # ==========================================================================
  # 私有辅助
  # ==========================================================================

  defp estimate_tokens(%State{token_usage: %{total_tokens: total}})
       when is_integer(total) and total > 0 do
    total
  end

  defp estimate_tokens(%State{messages: messages}) do
    messages
    |> Enum.reduce(0, fn msg, acc ->
      acc + div(byte_size(msg.content || ""), @chars_per_token)
    end)
  end

  defp put_arguments(tool_call, new_args) when is_map(tool_call) do
    cond do
      Map.has_key?(tool_call, :arguments) -> %{tool_call | arguments: new_args}
      Map.has_key?(tool_call, "arguments") -> Map.put(tool_call, "arguments", new_args)
      true -> Map.put(tool_call, :arguments, new_args)
    end
  end
end