cmdc

defmodule CMDC.Agent.Repair do
  # 内部模块 — Agent 消息历史修复器。API 不保证稳定。
  @moduledoc false

  require Logger

  @dialyzer {:no_opaque, find_orphaned_calls: 3}

  # ==========================================================================
  # Layer 2 — ensure_tool_results
  # ==========================================================================

  @doc """
  确保顺序消息列表中每条 assistant tool_calls 后紧跟对应的 tool_result。

  接受**顺序**消息列表（旧消息在前），返回修复后的顺序列表。
  """
  @spec ensure_tool_results([CMDC.Message.t()]) :: [CMDC.Message.t()]
  def ensure_tool_results(chronological_messages) do
    valid_call_ids =
      Enum.reduce(chronological_messages, MapSet.new(), fn
        %{role: :assistant, tool_calls: tcs}, ids when is_list(tcs) and tcs != [] ->
          tcs
          |> Enum.map(&tool_call_id/1)
          |> Enum.filter(&is_binary/1)
          |> Enum.reduce(ids, &MapSet.put(&2, &1))

        _, ids ->
          ids
      end)

    do_ensure(chronological_messages, valid_call_ids, [])
  end

  # ==========================================================================
  # Layer 1 — find_orphaned_calls
  # ==========================================================================

  @doc """
  扫描**倒序**消息列表，返回所有无对应 `tool_result` 的 tool_call 的 call_id 列表。
  """
  @spec find_orphaned_calls([CMDC.Message.t()]) :: [String.t()]
  def find_orphaned_calls(messages) do
    find_orphaned_calls(messages, MapSet.new(), [])
  end

  # ==========================================================================
  # 私有辅助 — do_ensure（Layer 2）
  # ==========================================================================

  defp do_ensure([], _valid_ids, acc), do: Enum.reverse(acc)

  defp do_ensure([%{role: :tool_result, call_id: cid} | rest], valid_ids, acc) do
    if MapSet.member?(valid_ids, cid) do
      Logger.warning("[Repair] stripping out-of-place tool_result: #{cid}")
    else
      Logger.warning("[Repair] stripping orphaned tool_result (no matching tool_call): #{cid}")
    end

    do_ensure(rest, valid_ids, acc)
  end

  defp do_ensure(
         [%{role: :assistant, tool_calls: tcs} = msg | rest],
         valid_ids,
         acc
       )
       when is_list(tcs) and tcs != [] do
    expected_ids =
      tcs
      |> Enum.map(&tool_call_id/1)
      |> Enum.filter(&is_binary/1)
      |> Enum.uniq()

    %{matched: matched, rest: rest} = take_tool_results_for_ids(rest, MapSet.new(expected_ids))

    missing = Enum.reject(expected_ids, &Map.has_key?(matched, &1))

    if missing != [] do
      Logger.warning(
        "[Repair] injecting #{length(missing)} synthetic tool_results for orphaned calls"
      )
    end

    results =
      Enum.map(expected_ids, fn call_id ->
        Map.get(matched, call_id) ||
          CMDC.Message.tool_result(call_id, "[Error: tool result missing]", true)
      end)

    new_acc = Enum.reverse([msg | results]) ++ acc
    do_ensure(rest, valid_ids, new_acc)
  end

  defp do_ensure([msg | rest], valid_ids, acc) do
    do_ensure(rest, valid_ids, [msg | acc])
  end

  # ==========================================================================
  # 私有辅助 — take_tool_results_for_ids
  # ==========================================================================

  @spec take_tool_results_for_ids([CMDC.Message.t()], MapSet.t()) :: %{
          matched: %{String.t() => CMDC.Message.t()},
          rest: [CMDC.Message.t()]
        }
  defp take_tool_results_for_ids(messages, expected_ids) do
    {matched, kept_rev} =
      Enum.reduce(messages, {%{}, []}, fn
        %{role: :tool_result, call_id: cid} = msg, {found, kept} ->
          cond do
            is_binary(cid) and MapSet.member?(expected_ids, cid) and not Map.has_key?(found, cid) ->
              {Map.put(found, cid, msg), kept}

            is_binary(cid) and MapSet.member?(expected_ids, cid) ->
              Logger.warning("[Repair] dropping duplicate tool_result for call_id: #{cid}")
              {found, kept}

            true ->
              {found, [msg | kept]}
          end

        msg, {found, kept} ->
          {found, [msg | kept]}
      end)

    %{matched: matched, rest: Enum.reverse(kept_rev)}
  end

  # ==========================================================================
  # 私有辅助 — find_orphaned_calls（Layer 1）
  # ==========================================================================

  defp find_orphaned_calls([], _result_ids, acc), do: acc

  defp find_orphaned_calls([%{role: :tool_result, call_id: cid} | rest], result_ids, acc) do
    find_orphaned_calls(rest, MapSet.put(result_ids, cid), acc)
  end

  defp find_orphaned_calls([%{role: :assistant, tool_calls: tcs} | rest], result_ids, acc)
       when is_list(tcs) and tcs != [] do
    orphans =
      tcs
      |> tool_call_ids()
      |> Enum.reject(&MapSet.member?(result_ids, &1))

    find_orphaned_calls(rest, result_ids, acc ++ orphans)
  end

  defp find_orphaned_calls([_ | rest], result_ids, acc) do
    find_orphaned_calls(rest, result_ids, acc)
  end

  # ==========================================================================
  # 私有辅助 — call_id 提取
  # ==========================================================================

  defp tool_call_id(%{call_id: call_id}) when is_binary(call_id), do: call_id
  defp tool_call_id(%{"call_id" => call_id}) when is_binary(call_id), do: call_id
  defp tool_call_id(_), do: nil

  defp tool_call_ids(tool_calls) do
    tool_calls
    |> Enum.reduce([], fn tc, ids ->
      case tool_call_id(tc) do
        call_id when is_binary(call_id) -> [call_id | ids]
        _ -> ids
      end
    end)
    |> Enum.reverse()
  end
end