opal

defmodule Opal.Session.Compaction do
  @moduledoc """
  Context window compaction following pi's approach.

  Summarizes older messages using the agent's LLM, producing a structured
  summary that preserves goals, progress, decisions, and file operations.
  Falls back to truncation if no agent is available.

  ## How it works

  1. Walk backwards from the newest message, estimating tokens, until
     `keep_recent_tokens` (default 20k) is reached — this is the cut point.
  2. Cut at a turn boundary (user message), never mid-turn.
  3. Serialize messages before the cut point into a text transcript.
  4. Ask the LLM to produce a structured summary.
  5. Replace old messages with a single summary message.

  ## Usage

      Opal.Session.Compaction.compact(session, agent: agent_pid)
      Opal.Session.Compaction.compact(session, strategy: :truncate)
  """

  require Logger

  @keep_recent_tokens 20_000
  @chars_per_token 4

  @summary_prompt """
  Summarize the following conversation transcript. Produce a structured summary using this exact format:

  ## Goal
  [What the user is trying to accomplish]

  ## Constraints & Preferences
  - [Requirements mentioned by user]

  ## Progress
  ### Done
  - [x] [Completed tasks]

  ### In Progress
  - [ ] [Current work]

  ## Key Decisions
  - **[Decision]**: [Rationale]

  ## Next Steps
  1. [What should happen next]

  ## Critical Context
  - [Any data or state needed to continue]

  <read-files>
  [list file paths that were read, one per line]
  </read-files>

  <modified-files>
  [list file paths that were created, written, or edited, one per line]
  </modified-files>

  Be concise. Focus on preserving information needed to continue the work.
  Do not include anything outside this format.

  ---

  TRANSCRIPT:

  """

  @doc """
  Compacts old messages in the session.

  ## Options

    * `:agent` — Agent pid for LLM summarization (calls get_state to get provider/model)
    * `:provider` — Provider module (alternative to `:agent`, avoids GenServer call)
    * `:model` — Model struct (required when `:provider` is given)
    * `:strategy` — `:summarize` (default if provider available) or `:truncate`
    * `:keep_recent_tokens` — tokens to keep uncompacted (default: #{@keep_recent_tokens})
    * `:instructions` — optional focus instructions for the summary
  """
  @spec compact(GenServer.server(), keyword()) :: :ok | {:error, term()}
  def compact(session, opts \\ []) do
    {provider, model} = resolve_provider(opts)
    has_provider = provider != nil
    strategy = Keyword.get(opts, :strategy, if(has_provider, do: :summarize, else: :truncate))
    keep_tokens = Keyword.get(opts, :keep_recent_tokens, @keep_recent_tokens)
    instructions = Keyword.get(opts, :instructions)

    path = Opal.Session.get_path(session)

    case find_cut_point(path, keep_tokens) do
      nil ->
        :ok

      cut_idx ->
        to_compact = Enum.slice(path, 0, cut_idx)
        ids_to_remove = Enum.map(to_compact, & &1.id)

        summary_content = build_summary(to_compact, strategy, {provider, model}, instructions)

        summary_msg = %Opal.Message{
          id: generate_id(),
          role: :user,
          content: "[Conversation summary — older messages were compacted]\n\n#{summary_content}",
          parent_id: nil
        }

        Logger.debug("Compacting #{length(to_compact)} messages, keeping #{length(path) - cut_idx}")
        Opal.Session.replace_path_segment(session, ids_to_remove, summary_msg)
    end
  end

  # Resolve provider/model from opts — either explicit or via agent pid
  defp resolve_provider(opts) do
    case {Keyword.get(opts, :provider), Keyword.get(opts, :model)} do
      {p, m} when p != nil and m != nil -> {p, m}
      _ ->
        case Keyword.get(opts, :agent) do
          nil -> {nil, nil}
          agent ->
            state = Opal.Agent.get_state(agent)
            {state.provider, state.model}
        end
    end
  end

  # Walk backwards from the end, accumulating estimated tokens.
  # Return the index of the first message to keep (cut point).
  # Always cut at a user message boundary to avoid splitting turns.
  defp find_cut_point(path, keep_tokens) do
    total = length(path)
    keep_chars = keep_tokens * @chars_per_token

    {_, _acc_chars, cut_idx} =
      path
      |> Enum.reverse()
      |> Enum.with_index()
      |> Enum.reduce({false, 0, nil}, fn {msg, rev_idx}, {found, chars, cut} ->
        msg_chars = estimate_chars(msg)
        new_chars = chars + msg_chars

        if found do
          {true, new_chars, cut}
        else
          if new_chars >= keep_chars do
            # Find the nearest user message boundary at or before this point
            real_idx = total - 1 - rev_idx
            boundary = find_turn_boundary(path, real_idx)
            {true, new_chars, boundary}
          else
            {false, new_chars, cut}
          end
        end
      end)

    # Only compact if we found a cut point with at least one message to remove
    if cut_idx && cut_idx >= 1, do: cut_idx, else: nil
  end

  # Find the nearest turn boundary (user message) at or after the given index.
  defp find_turn_boundary(path, idx) do
    path
    |> Enum.with_index()
    |> Enum.drop(idx)
    |> Enum.find_value(fn {msg, i} ->
      if msg.role == :user, do: i
    end) || idx
  end

  defp estimate_chars(%{content: nil}), do: 20
  defp estimate_chars(%{content: c}) when is_binary(c), do: byte_size(c) + 20
  defp estimate_chars(_), do: 20

  # --- Summary generation ---

  defp build_summary(messages, :truncate, _provider_model, _instructions) do
    file_ops = extract_file_ops(messages)
    count = length(messages)
    roles = messages |> Enum.map(& &1.role) |> Enum.frequencies()
    role_str = roles |> Enum.map(fn {r, n} -> "#{n} #{r}" end) |> Enum.join(", ")

    read = if file_ops.read != [], do: "\n\n<read-files>\n#{Enum.join(file_ops.read, "\n")}\n</read-files>", else: ""
    modified = if file_ops.modified != [], do: "\n\n<modified-files>\n#{Enum.join(file_ops.modified, "\n")}\n</modified-files>", else: ""

    "[Compacted #{count} messages: #{role_str}]#{read}#{modified}"
  end

  defp build_summary(messages, :summarize, {provider, model}, instructions) do
    transcript = serialize_conversation(messages)

    prompt =
      if instructions do
        @summary_prompt <> transcript <> "\n\nAdditional focus instructions: #{instructions}"
      else
        @summary_prompt <> transcript
      end

    case summarize_with_provider(provider, model, prompt) do
      {:ok, summary} ->
        summary

      {:error, _reason} ->
        Logger.warning("LLM summarization failed, falling back to truncation")
        build_summary(messages, :truncate, nil, nil)
    end
  end

  defp summarize_with_provider(provider, model, prompt) do
    summary_messages = [
      %Opal.Message{id: "sys", role: :system, content: "You are a conversation summarizer. Produce only the structured summary, nothing else."},
      Opal.Message.user(prompt)
    ]

    case provider.stream(model, summary_messages, []) do
      {:ok, resp} ->
        text = collect_stream_text(resp, provider, "")
        if text != "", do: {:ok, String.trim(text)}, else: {:error, :empty}

      {:error, reason} ->
        {:error, reason}
    end
  end

  # Synchronously collects all text from a streaming response.
  defp collect_stream_text(resp, provider, acc) do
    receive do
      message ->
        case Req.parse_message(resp, message) do
          {:ok, chunks} when is_list(chunks) ->
            new_acc =
              Enum.reduce(chunks, acc, fn
                {:data, data}, text_acc ->
                  binary = IO.iodata_to_binary(data)

                  binary
                  |> String.split("\n", trim: true)
                  |> Enum.reduce(text_acc, fn
                    "data: [DONE]", inner -> inner
                    "data: " <> json, inner ->
                      events = provider.parse_stream_event(json)
                      Enum.reduce(events, inner, fn
                        {:text_delta, delta}, t -> t <> delta
                        _, t -> t
                      end)
                    _, inner -> inner
                  end)

                :done, text_acc -> text_acc
                _, text_acc -> text_acc
              end)

            if :done in chunks, do: new_acc, else: collect_stream_text(resp, provider, new_acc)

          :unknown ->
            collect_stream_text(resp, provider, acc)
        end
    after
      30_000 -> acc
    end
  end

  @doc """
  Serializes messages into a text transcript for summarization.
  """
  @spec serialize_conversation([Opal.Message.t()]) :: String.t()
  def serialize_conversation(messages) do
    messages
    |> Enum.map(fn msg ->
      case msg.role do
        :user ->
          "[User]: #{msg.content || ""}"

        :assistant ->
          lines = ["[Assistant]: #{msg.content || ""}"]

          tool_lines =
            case msg.tool_calls do
              calls when is_list(calls) and calls != [] ->
                tools = Enum.map_join(calls, "; ", fn tc ->
                  args = tc.arguments |> Jason.encode!() |> String.slice(0, 200)
                  "#{tc.name}(#{args})"
                end)
                ["[Assistant tool calls]: #{tools}"]

              _ ->
                []
            end

          Enum.join(lines ++ tool_lines, "\n")

        :tool_result ->
          output = String.slice(msg.content || "", 0, 500)
          "[Tool result (#{msg.name || msg.call_id})]: #{output}"

        :system ->
          "[System]: #{msg.content || ""}"

        _ ->
          "[#{msg.role}]: #{msg.content || ""}"
      end
    end)
    |> Enum.join("\n")
  end

  @doc """
  Extracts file read/write operations from tool calls in messages.
  """
  @spec extract_file_ops([Opal.Message.t()]) :: %{read: [String.t()], modified: [String.t()]}
  def extract_file_ops(messages) do
    Enum.reduce(messages, %{read: [], modified: []}, fn msg, acc ->
      case msg do
        %{role: :assistant, tool_calls: calls} when is_list(calls) ->
          Enum.reduce(calls, acc, fn tc, inner ->
            case tc.name do
              "read" ->
                path = tc.arguments["path"]
                if path, do: %{inner | read: [path | inner.read]}, else: inner

              name when name in ["write", "edit"] ->
                path = tc.arguments["path"]
                if path, do: %{inner | modified: [path | inner.modified]}, else: inner

              _ ->
                inner
            end
          end)

        _ ->
          acc
      end
    end)
    |> then(fn ops ->
      %{read: Enum.uniq(Enum.reverse(ops.read)), modified: Enum.uniq(Enum.reverse(ops.modified))}
    end)
  end

  defp generate_id do
    :crypto.strong_rand_bytes(8) |> Base.encode16(case: :lower)
  end
end