cmdc

defmodule CMDC.Agent.Compactor do
  @moduledoc """
  上下文压缩引擎 — 对标 DeepAgents `SummarizationMiddleware`。

  ## 职责

  - 监控上下文 token 使用量，在超过阈值时自动压缩旧消息
  - 将被驱逐的消息 offload 到文件，保留完整历史
  - ContextOverflow 时作为 fallback 强制压缩
  - 生成摘要消息替换旧历史

  ## Token 计数策略

  采用 **Provider usage 为主 + 字符估算 fallback**：

  1. 优先使用 `State.token_usage.total_tokens`（Provider 返回的真实值）
  2. Fallback：消息文本 `byte_size / 4` 粗估

  ## 触发条件（二选一）

  - `:token` 模式 — `total_tokens >= trigger_tokens`（默认 170_000）
  - `:fraction` 模式 — `total_tokens >= context_window * trigger_fraction`（默认 0.85）

  ## 保留策略

  - `:messages` 模式 — 保留最近 N 条消息（默认 6）
  - `:fraction` 模式 — 保留最近占 context_window 的 fraction

  ## 使用方式

  不直接调用，由 Agent `run_turn` 在 `before_request` 后自动检查：

      case Compactor.maybe_compact(state) do
        {:compacted, state} -> run_turn_with(state)
        {:skip, state} -> run_turn_with(state)
      end

  或在 ContextOverflow 时强制调用：

      {:compacted, state} = Compactor.force_compact(state)
  """

  require Logger
  alias CMDC.Agent.{Emitter, State}

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

  @default_trigger_tokens 170_000
  @default_keep_messages 6
  @default_keep_fraction 0.10
  @default_offload_dir "/conversation_history"
  @chars_per_token 4

  @summary_prompt """
  You are summarizing a conversation between a user and an AI assistant.
  Produce a concise summary that preserves:
  - Key decisions made
  - Important file paths, names, and code snippets mentioned
  - Current task status and next steps
  - Any errors encountered and their resolutions

  Be factual and concise. Use bullet points. Do not add commentary.
  """

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

  @typedoc "压缩触发配置。"
  @type trigger :: {:tokens, pos_integer()} | {:fraction, float()}

  @typedoc "消息保留配置。"
  @type keep :: {:messages, pos_integer()} | {:fraction, float()}

  @typedoc "压缩配置。"
  @type compact_opts :: %{
          trigger: trigger(),
          keep: keep(),
          offload_dir: String.t(),
          summary_prompt: String.t()
        }

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

  @doc """
  检查是否需要压缩，需要时自动执行。

  返回：
  - `{:compacted, state}` — 已执行压缩，state 中消息已替换
  - `{:skip, state}` — 未达到阈值，不需要压缩
  """
  @spec maybe_compact(State.t()) :: {:compacted, State.t()} | {:skip, State.t()}
  def maybe_compact(%State{} = state) do
    opts = compact_config(state)
    total_tokens = estimate_tokens(state)

    if should_compact?(total_tokens, opts, state) do
      do_compact(state, opts)
    else
      {:skip, state}
    end
  end

  @doc """
  仅检查是否达到压缩阈值（不执行压缩）。

  供 Agent 层在 `:before_compact` Plugin Pipeline 之前判断是否触发。
  """
  @spec should_compact?(State.t()) :: boolean()
  def should_compact?(%State{} = state) do
    opts = compact_config(state)
    total_tokens = estimate_tokens(state)
    should_compact?(total_tokens, opts, state)
  end

  @doc """
  强制执行压缩（ContextOverflow fallback）。

  忽略阈值检查，直接压缩。如果消息太少无法分割，返回原 state。
  """
  @spec force_compact(State.t()) :: {:compacted, State.t()} | {:skip, State.t()}
  def force_compact(%State{} = state) do
    opts = compact_config(state)
    do_compact(state, opts)
  end

  @doc """
  估算当前上下文 token 数量。

  优先使用 Provider 返回的 usage，fallback 到字符估算。
  """
  @spec estimate_tokens(State.t()) :: non_neg_integer()
  def estimate_tokens(%State{token_usage: %{total_tokens: total}})
      when is_integer(total) and total > 0 do
    total
  end

  def estimate_tokens(%State{messages: messages}) do
    estimate_tokens_from_messages(messages)
  end

  @doc """
  从消息列表估算 token 数（字符数 / #{@chars_per_token} 粗估）。
  """
  @spec estimate_tokens_from_messages([CMDC.Message.t()]) :: non_neg_integer()
  def estimate_tokens_from_messages(messages) when is_list(messages) do
    messages
    |> Enum.reduce(0, fn msg, acc ->
      content_size = byte_size(msg.content || "")
      thinking_size = byte_size(msg.thinking || "")
      tool_calls_size = estimate_tool_calls_tokens(msg.tool_calls)
      acc + div(content_size + thinking_size + tool_calls_size, @chars_per_token)
    end)
  end

  @doc """
  读取压缩配置。从 `state.config` 中获取，缺省使用默认值。

  支持的 config 键：
  - `:compact_trigger` — `{:tokens, n}` 或 `{:fraction, f}`
  - `:compact_keep` — `{:messages, n}` 或 `{:fraction, f}`
  - `:compact_offload_dir` — offload 文件目录前缀
  - `:compact_summary_prompt` — 摘要生成提示词
  """
  @spec compact_config(State.t()) :: compact_opts()
  def compact_config(%State{config: config}) when is_map(config) do
    %{
      trigger: Map.get(config, :compact_trigger, {:tokens, @default_trigger_tokens}),
      keep: Map.get(config, :compact_keep, {:messages, @default_keep_messages}),
      offload_dir: Map.get(config, :compact_offload_dir, @default_offload_dir),
      summary_prompt: Map.get(config, :compact_summary_prompt, @summary_prompt)
    }
  end

  def compact_config(_state) do
    %{
      trigger: {:tokens, @default_trigger_tokens},
      keep: {:messages, @default_keep_messages},
      offload_dir: @default_offload_dir,
      summary_prompt: @summary_prompt
    }
  end

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

  defp should_compact?(total_tokens, opts, state) do
    case opts.trigger do
      {:tokens, threshold} ->
        total_tokens >= threshold

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

        if context_window > 0 do
          total_tokens >= trunc(context_window * fraction)
        else
          total_tokens >= trunc(@default_trigger_tokens * fraction)
        end
    end
  end

  # ==========================================================================
  # 私有 — 压缩执行
  # ==========================================================================

  defp do_compact(state, opts) do
    chronological = Enum.reverse(state.messages)

    case partition_messages(chronological, opts, state) do
      {[], _keep} ->
        {:skip, state}

      {to_summarize, to_keep} ->
        Logger.info(
          "[Compactor] compacting: #{length(to_summarize)} msgs → summary, keeping #{length(to_keep)} msgs"
        )

        broadcast(state, {:compact_start, state.session_id})

        offload_path = offload_messages(state, to_summarize, opts)
        summary_text = generate_summary(state, to_summarize, opts)
        summary_msg = build_summary_message(summary_text, offload_path)

        new_messages = [summary_msg | to_keep] |> Enum.reverse()

        state = %{state | messages: new_messages}

        broadcast(
          state,
          {:compact_end, state.session_id, length(to_summarize)}
        )

        {:compacted, state}
    end
  end

  # ==========================================================================
  # 私有 — 消息分割
  # ==========================================================================

  @doc false
  @spec partition_messages([CMDC.Message.t()], compact_opts(), State.t()) ::
          {[CMDC.Message.t()], [CMDC.Message.t()]}
  def partition_messages(chronological, opts, state) do
    keep_count = determine_keep_count(chronological, opts, state)
    total = length(chronological)

    if keep_count >= total do
      {[], chronological}
    else
      cutoff = total - keep_count
      {to_summarize, to_keep} = Enum.split(chronological, cutoff)

      to_keep = ensure_valid_start(to_keep, to_summarize)
      to_summarize = Enum.take(chronological, total - length(to_keep))

      {to_summarize, to_keep}
    end
  end

  defp determine_keep_count(messages, opts, state) do
    case opts.keep do
      {:messages, n} ->
        n

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

        target_tokens =
          if context_window > 0 do
            trunc(context_window * fraction)
          else
            trunc(@default_trigger_tokens * @default_keep_fraction)
          end

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

  defp count_messages_for_tokens(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

  # tool_result 消息不能作为保留区首条——需要把前面的 assistant（含 tool_call）也拉进来
  defp ensure_valid_start(to_keep, to_summarize) do
    case to_keep do
      [%{role: :tool_result} | _] ->
        pull_assistant_back(Enum.reverse(to_summarize), to_keep)

      _ ->
        to_keep
    end
  end

  # 从 to_summarize 尾部拉回消息，直到 to_keep 不再以 tool_result 开头
  defp pull_assistant_back([], keep), do: keep

  defp pull_assistant_back(summarize_reversed, keep) do
    [last | rest] = summarize_reversed
    new_keep = [last | keep]

    case new_keep do
      [%{role: :tool_result} | _] -> pull_assistant_back(rest, new_keep)
      _ -> new_keep
    end
  end

  # ==========================================================================
  # 私有 — 摘要生成
  # ==========================================================================

  defp generate_summary(state, messages, opts) do
    summary_fn = Map.get(state.config, :compact_summary_fn)

    if is_function(summary_fn, 2) do
      case summary_fn.(messages, opts.summary_prompt) do
        {:ok, text} -> text
        {:error, _} -> fallback_summary(messages)
        text when is_binary(text) -> text
      end
    else
      fallback_summary(messages)
    end
  end

  defp fallback_summary(messages) do
    lines =
      Enum.map_join(messages, "\n", fn msg ->
        role = msg.role |> to_string() |> String.upcase()
        content = truncate_content(msg.content || "", 200)
        "#{role}: #{content}"
      end)

    "## Conversation Summary (#{length(messages)} messages)\n\n#{lines}"
  end

  defp truncate_content(text, max_len) do
    if byte_size(text) > max_len do
      String.slice(text, 0, max_len) <> "..."
    else
      text
    end
  end

  # ==========================================================================
  # 私有 — Offload（写入文件）
  # ==========================================================================

  defp offload_messages(state, messages, opts) do
    sandbox = state.sandbox
    working_dir = state.working_dir

    path = Path.join(opts.offload_dir, "#{state.session_id}.md")

    content = format_offload_content(messages)

    result =
      if sandbox do
        sandbox.write_file(path, content, working_dir: working_dir)
      else
        full_path = Path.join(working_dir, path)
        File.mkdir_p!(Path.dirname(full_path))

        case File.write(full_path, content) do
          :ok -> :ok
          {:error, reason} -> {:error, inspect(reason)}
        end
      end

    case result do
      :ok ->
        Logger.debug("[Compactor] offloaded #{length(messages)} messages to #{path}")
        path

      {:error, reason} ->
        Logger.warning("[Compactor] offload failed: #{reason}")
        nil
    end
  end

  defp format_offload_content(messages) do
    timestamp = DateTime.utc_now() |> DateTime.to_iso8601()

    header = "## Compacted at #{timestamp}\n\n"

    body =
      Enum.map_join(messages, "\n", fn msg ->
        role = msg.role |> to_string()
        content = msg.content || ""
        "### #{role}\n\n#{content}\n"
      end)

    header <> body
  end

  # ==========================================================================
  # 私有 — 摘要消息构建
  # ==========================================================================

  defp build_summary_message(summary_text, offload_path) do
    content =
      if offload_path do
        "You are in the middle of a conversation that has been summarized.\n\n" <>
          "The full conversation history has been saved to #{offload_path} " <>
          "should you need to refer back to it for details.\n\n" <>
          "A condensed summary follows:\n\n" <>
          "<summary>\n#{summary_text}\n</summary>"
      else
        "Here is a summary of the conversation to date:\n\n#{summary_text}"
      end

    CMDC.Message.user(content)
  end

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

  defp estimate_tool_calls_tokens(nil), do: 0
  defp estimate_tool_calls_tokens([]), do: 0

  defp estimate_tool_calls_tokens(tool_calls) when is_list(tool_calls) do
    Enum.reduce(tool_calls, 0, fn tc, acc ->
      name_size = byte_size(tc[:name] || "")
      args_size = tc[:arguments] |> inspect() |> byte_size()
      acc + name_size + args_size
    end)
  end

  defp broadcast(%State{} = state, event) do
    Emitter.broadcast(state, event)
  end
end