cmdc

defmodule CMDC.Agent.Stream do
  # 内部模块 — StreamChunk 事件处理与状态更新。API 不保证稳定。
  @moduledoc false

  require Logger
  alias CMDC.Agent.State
  alias CMDC.EventBus

  @dialyzer [{:nowarn_function, broadcast: 2}]

  # ==========================================================================
  # StreamChunk 处理入口
  # ==========================================================================

  @doc """
  处理单个 `ReqLLM.StreamChunk`，更新 Agent State。

  是 gen_statem streaming 状态收到 `{:cmdc_stream_chunk, chunk}` 消息后的核心处理函数。
  通过 chunk.type 分发到不同的处理逻辑。
  """
  @spec handle_chunk(term(), State.t()) :: State.t()
  def handle_chunk(%{type: :content, text: text}, state)
      when is_binary(text) do
    state = ensure_message_started(state)
    {clean, state} = extract_status_tags(text, state)
    {clean, state} = extract_title_tag(clean, state)

    if clean != "", do: broadcast(state, {:message_delta, %{delta: clean}})
    %{state | current_text: state.current_text <> clean}
  end

  def handle_chunk(%{type: :thinking, text: text}, state)
      when is_binary(text) do
    state =
      if is_nil(state.current_thinking) do
        broadcast(state, :thinking_start)
        %{state | current_thinking: ""}
      else
        state
      end

    broadcast(state, {:thinking_delta, %{delta: text}})
    %{state | current_thinking: state.current_thinking <> text}
  end

  def handle_chunk(%{type: :tool_call, name: name, arguments: arguments, metadata: meta}, state)
      when is_binary(name) do
    meta = meta || %{}

    tool_call = %{
      name: name,
      arguments: arguments || %{},
      call_id: Map.get(meta, :id),
      block_index: Map.get(meta, :index)
    }

    %{state | current_tool_calls: upsert_tool_call(state.current_tool_calls, tool_call)}
  end

  # :tool_call chunk 有时缺少 metadata 字段
  def handle_chunk(%{type: :tool_call, name: name, arguments: arguments} = chunk, state)
      when is_binary(name) and not is_map_key(chunk, :metadata) do
    tool_call = %{
      name: name,
      arguments: arguments || %{},
      call_id: nil,
      block_index: nil
    }

    %{state | current_tool_calls: upsert_tool_call(state.current_tool_calls, tool_call)}
  end

  def handle_chunk(
        %{
          type: :meta,
          metadata: %{tool_call_args: %{index: index, fragment: fragment}} = metadata
        },
        state
      )
      when is_integer(index) and is_binary(fragment) do
    buf = Map.get(state.tool_call_arg_buffers, index, [])

    state = %{
      state
      | tool_call_arg_buffers: Map.put(state.tool_call_arg_buffers, index, [buf, fragment])
    }

    maybe_update_usage(state, metadata)
  end

  def handle_chunk(%{type: :meta, metadata: metadata}, state)
      when is_map(metadata) do
    maybe_update_usage(state, metadata)
  end

  # 前向兼容：忽略未知 chunk 类型
  def handle_chunk(_unknown, state), do: state

  # ==========================================================================
  # XML 标签提取
  # ==========================================================================

  @doc """
  通用 XML 标签提取器，处理跨 chunk 的流式 delta。

  从 `delta` 中剔除 `<tag_name>…</tag_name>`，发现完整标签时调用 `on_match`。
  不完整的标签缓冲到 `state.tag_buffers[tag_name]`，等待下一个 chunk。
  """
  @spec extract_xml_tag(String.t(), atom(), State.t(), (String.t(), State.t() -> State.t())) ::
          {String.t(), State.t()}
  def extract_xml_tag(delta, tag_name, state, on_match) do
    buf = Map.get(state.tag_buffers, tag_name, "")
    text = buf <> delta
    close_tag = "</#{tag_name}>"
    tag_regex = ~r/<#{tag_name}>(.*?)<\/#{tag_name}>/s

    case Regex.run(tag_regex, text) do
      [full_match, inner_text] ->
        state = on_match.(String.trim(inner_text), state)
        clean = String.replace(text, full_match, "", global: false)

        {more_clean, state} =
          extract_xml_tag("", tag_name, clear_tag_buffer(state, tag_name), on_match)

        {clean <> more_clean, state}

      nil ->
        open_tag = "<#{tag_name}>"

        cond do
          String.contains?(text, open_tag) and not String.contains?(text, close_tag) ->
            [before | _] = String.split(text, open_tag, parts: 2)
            rest = String.slice(text, String.length(before)..-1//1)
            {before, put_tag_buffer(state, tag_name, rest)}

          (len = partial_tag_suffix_length(text, open_tag)) > 0 ->
            split_at = String.length(text) - len

            {String.slice(text, 0, split_at),
             put_tag_buffer(state, tag_name, String.slice(text, split_at..-1//1))}

          true ->
            {text, clear_tag_buffer(state, tag_name)}
        end
    end
  end

  @doc "从流式 delta 中提取 `<status>…</status>` 标签，广播 `{:status_update, text}`。"
  @spec extract_status_tags(String.t(), State.t()) :: {String.t(), State.t()}
  def extract_status_tags(delta, state) do
    extract_xml_tag(delta, :status, state, fn text, st ->
      broadcast(st, {:status_update, text})
      st
    end)
  end

  @doc "从流式 delta 中提取 `<title>…</title>` 标签，广播 `{:title_generated, title}`（限 60 字符）。"
  @spec extract_title_tag(String.t(), State.t()) :: {String.t(), State.t()}
  def extract_title_tag(delta, state) do
    extract_xml_tag(delta, :title, state, fn text, st ->
      title = String.slice(text, 0, 60)
      if title != "", do: broadcast(st, {:title_generated, title})
      st
    end)
  end

  @doc "返回 `text` 结尾匹配 `<status>` 前缀的字符数。"
  @spec partial_tag_length(String.t()) :: non_neg_integer()
  def partial_tag_length(text) do
    partial_tag_suffix_length(text, "<status>")
  end

  @doc "返回 `text` 结尾匹配 `<title>` 前缀的字符数。"
  @spec partial_title_tag_length(String.t()) :: non_neg_integer()
  def partial_title_tag_length(text) do
    partial_tag_suffix_length(text, "<title>")
  end

  # ==========================================================================
  # 私有辅助 — 消息开始
  # ==========================================================================

  defp ensure_message_started(%{message_started: true} = state), do: state

  defp ensure_message_started(state) do
    broadcast(state, :message_start)
    %{state | message_started: true}
  end

  # ==========================================================================
  # 私有辅助 — Usage 更新
  # ==========================================================================

  defp maybe_update_usage(state, %{usage: usage}) when is_map(usage) do
    State.update_token_usage(state, normalize_usage(usage))
  end

  defp maybe_update_usage(state, _metadata), do: state

  # 把 Provider raw usage 归一化为内部 token map（State 内部为兼容性仍用 map，含
  # 上下文窗口元数据；公开事件统一用 %CMDC.TokenUsage{}）。
  defp normalize_usage(usage) do
    struct = CMDC.TokenUsage.from_raw(usage)

    base = %{
      prompt_tokens: struct.prompt_tokens,
      completion_tokens: struct.completion_tokens,
      total_tokens: struct.total_tokens
    }

    if is_integer(struct.cached_tokens) do
      Map.put(base, :cached_tokens, struct.cached_tokens)
    else
      base
    end
  end

  # ==========================================================================
  # 私有辅助 — Tool Call 组装
  # ==========================================================================

  defp upsert_tool_call(tool_calls, %{name: name} = tool_call) do
    case find_last_index(tool_calls, &(&1[:name] == name and is_nil(&1[:call_id]))) do
      nil ->
        tool_calls ++ [tool_call]

      idx ->
        List.update_at(tool_calls, idx, fn existing ->
          existing
          |> Map.put(:arguments, merge_arguments(existing[:arguments], tool_call[:arguments]))
          |> maybe_put(:call_id, tool_call[:call_id])
        end)
    end
  end

  defp merge_arguments(nil, incoming), do: incoming || %{}
  defp merge_arguments(existing, nil), do: existing
  defp merge_arguments(existing, incoming) when map_size(existing) == 0, do: incoming
  defp merge_arguments(existing, incoming), do: Map.merge(existing, incoming)

  defp maybe_put(map, _key, nil), do: map
  defp maybe_put(map, key, value), do: Map.put(map, key, value)

  defp find_last_index(list, predicate) do
    list
    |> Enum.with_index()
    |> Enum.reduce(nil, fn {item, idx}, acc ->
      if predicate.(item), do: idx, else: acc
    end)
  end

  # ==========================================================================
  # 私有辅助 — Tag Buffer
  # ==========================================================================

  defp put_tag_buffer(state, tag_name, value) do
    %{state | tag_buffers: Map.put(state.tag_buffers, tag_name, value)}
  end

  defp clear_tag_buffer(state, tag_name) do
    %{state | tag_buffers: Map.put(state.tag_buffers, tag_name, "")}
  end

  defp partial_tag_suffix_length(text, open_tag) do
    tag_without_close = String.trim_trailing(open_tag, ">")
    len = String.length(tag_without_close)

    Enum.find_value(len..1//-1, 0, fn n ->
      suffix = String.slice(tag_without_close, 0, n)
      if String.ends_with?(text, suffix), do: n
    end)
  end

  # ==========================================================================
  # 私有辅助 — 事件广播
  # ==========================================================================

  defp broadcast(%{session_id: session_id}, event) do
    EventBus.broadcast(session_id, event)
  end
end