eai

defmodule Eai.Hook.ContextMeter do
  @moduledoc """
  Observability hook: prints context size after every LLM HTTP response.

  Fires on :llm_post for each HTTP response (including tool-call intermediate
  steps), so the user can watch context grow in real-time and decide when
  to interrupt/summarize.

  Output examples:

      📊 [tool] context: 120.4 KB (~30.1K tokens), 18 messages
      📊 [final] context: 206.4 KB (~52.8K tokens), 22 messages

  Token estimate uses 4 bytes/token (English avg). This is a rough guide
  for the user to decide when to summarize/export — not an exact bill.
  """

  use Eai.Hook, priority: 30

  alias Eai.Message

  @impl true
  def interest(:llm_post, "LLM_REQUEST", _payload), do: true
  def interest(_event, _tool, _payload), do: false

  @impl true
  def verdict(:llm_post, _tool, _payload, result) do
    case result do
      {:ok, _reply, history} when is_list(history) and history != [] ->
        print_meter(history)

      _ ->
        :ok
    end
  end

  defp print_meter(history) do
    bytes = estimate_bytes(history)
    msg_count = length(history)
    token_est = div(bytes, 4)
    tag = context_tag(history)

    IO.puts(
      "\n📊 [#{tag}] context: #{format_bytes(bytes)} (~#{format_tokens(token_est)} tokens), #{msg_count} messages"
    )
  end

  # 标注当前上下文状态：工具调用中间步还是最终回复
  defp context_tag(history) do
    last = List.last(history)

    case last.role do
      :assistant ->
        if Message.has_tool_uses?(last), do: "tool", else: "final"

      _ ->
        "step"
    end
  end

  # 估算消息列表的序列化字节数。
  defp estimate_bytes(history) do
    history
    |> :erlang.term_to_binary()
    |> byte_size()
  end

  defp format_bytes(bytes) when bytes >= 1024 * 1024 do
    :erlang.float_to_binary(bytes / (1024 * 1024), decimals: 1) <> " MB"
  end

  defp format_bytes(bytes) when bytes >= 1024 do
    :erlang.float_to_binary(bytes / 1024, decimals: 1) <> " KB"
  end

  defp format_bytes(bytes), do: "#{bytes} B"

  defp format_tokens(tokens) when tokens >= 1000 do
    :erlang.float_to_binary(tokens / 1000, decimals: 1) <> "K"
  end

  defp format_tokens(tokens), do: "#{tokens}"
end