cmdc_test

defmodule CMDCTest.RAG.Assertions do
  @moduledoc """
  RAG / GraphRAG contract assertions。

  支持输入 Tool JSON、map 或 RAG plugin event payload。断言重点是 citation
  provenance、ACL source 泄漏、pipeline summary、grounding 与 graph evidence。
  """

  import ExUnit.Assertions

  @raw_text_keys ~w(text chunk_text content prompt chunks raw raw_results results)a
  @raw_text_string_keys Enum.map(@raw_text_keys, &Atom.to_string/1)

  @doc "断言结果至少包含指定数量的 citations,并返回 citations。"
  @spec assert_citations(String.t() | map(), keyword()) :: [map()]
  def assert_citations(result, opts \\ []) do
    map = normalize_result(result)
    citations = list_field(map, :citations)
    min = Keyword.get(opts, :min, 1)

    assert length(citations) >= min,
           "expected at least #{min} citations, got #{length(citations)}"

    Enum.each(citations, fn citation ->
      assert present?(field(citation, :document_id)) or present?(field(citation, :source_uri)),
             "citation must include document_id or source_uri: #{inspect(citation)}"
    end)

    maybe_assert_any(citations, :collection, Keyword.get(opts, :collection))
    maybe_assert_any(citations, :document_id, Keyword.get(opts, :document_id))
    maybe_assert_any(citations, :source_uri, Keyword.get(opts, :source_uri))

    citations
  end

  @doc "断言 citations 没有越权 collection。"
  @spec assert_no_unauthorized_sources(String.t() | map(), [String.t()], keyword()) :: :ok
  def assert_no_unauthorized_sources(result, allowed_collections, opts \\ []) do
    citations = assert_citations(result, Keyword.take(opts, [:min]))
    allowed = MapSet.new(Enum.map(allowed_collections, &to_string/1))

    unauthorized =
      Enum.reject(citations, fn citation ->
        collection = field(citation, :collection)
        is_nil(collection) or MapSet.member?(allowed, to_string(collection))
      end)

    assert unauthorized == [],
           "expected no unauthorized citation collections, got #{inspect(unauthorized)}"

    :ok
  end

  @doc "断言结果有 grounding score,并达到最小阈值。"
  @spec assert_grounding(String.t() | map(), number()) :: map()
  def assert_grounding(result, min_score \\ 0.0) do
    metadata = metadata(result)
    grounding = field(metadata, :grounding) || %{}
    score = numeric_field(grounding, :score) || numeric_field(metadata, :grounding_score)

    assert is_number(score), "expected grounding score in result metadata"
    assert score >= min_score, "expected grounding score >= #{min_score}, got #{score}"

    grounding
  end

  @doc "断言结果包含 pipeline run summary。"
  @spec assert_pipeline_summary(String.t() | map(), keyword()) :: map()
  def assert_pipeline_summary(result, opts \\ []) do
    metadata = metadata(result)
    summary = field(metadata, :pipeline_run_summary) || field(metadata, :pipeline)

    assert is_map(summary), "expected pipeline run summary in result metadata"

    case Keyword.get(opts, :preset_id) do
      nil -> :ok
      preset_id -> assert field(summary, :preset_id) == preset_id
    end

    min_steps = Keyword.get(opts, :min_steps, 1)
    steps = list_field(summary, :steps)

    assert length(steps) >= min_steps,
           "expected at least #{min_steps} pipeline steps, got #{length(steps)}"

    summary
  end

  @doc "断言结果包含 GraphRAG evidence。"
  @spec assert_graph_evidence(String.t() | map(), keyword()) :: map()
  def assert_graph_evidence(result, opts \\ []) do
    metadata = metadata(result)
    min_entities = Keyword.get(opts, :min_entities, 1)
    min_relationships = Keyword.get(opts, :min_relationships, 1)

    entities = list_field(metadata, :entity_support)
    relationships = list_field(metadata, :relationship_support)

    assert length(entities) >= min_entities,
           "expected at least #{min_entities} graph entity evidence items, got #{length(entities)}"

    assert length(relationships) >= min_relationships,
           "expected at least #{min_relationships} graph relationship evidence items, got #{length(relationships)}"

    %{
      entity_support: entities,
      relationship_support: relationships,
      path_support: list_field(metadata, :path_support),
      community_support: list_field(metadata, :community_support)
    }
  end

  @doc "断言 RAG 审计/进度 payload 没有泄漏正文、chunk 或 prompt。"
  @spec refute_raw_text_leaked(map() | tuple() | String.t()) :: :ok
  def refute_raw_text_leaked(payload) do
    payload = normalize_event_payload(payload)
    leaks = collect_raw_text_keys(payload, [])

    assert leaks == [],
           "expected RAG event payload not to include raw text/chunks/prompt, got #{inspect(leaks)}"

    :ok
  end

  @doc "断言 Phase 18 RAG plugin event 名称与 payload 子集。"
  @spec assert_rag_event(tuple(), atom(), keyword()) :: map()
  def assert_rag_event({:plugin_event, name, payload}, expected_name, opts \\ []) do
    assert name == expected_name

    expected_payload = Keyword.get(opts, :payload, %{})

    Enum.each(expected_payload, fn {key, value} ->
      assert field(payload, key) == value
    end)

    payload
  end

  defp normalize_result(result) when is_binary(result) do
    case Jason.decode(result) do
      {:ok, map} when is_map(map) -> map
      {:ok, other} -> flunk("expected JSON object, got #{inspect(other)}")
      {:error, error} -> flunk("failed to decode JSON result: #{inspect(error)}")
    end
  end

  defp normalize_result(%{} = result), do: result

  defp normalize_event_payload({:plugin_event, _name, payload}), do: payload
  defp normalize_event_payload(payload), do: normalize_result(payload)

  defp metadata(result), do: normalize_result(result) |> field(:metadata) || %{}

  defp field(%{} = map, key) when is_atom(key) do
    Map.get(map, key) || Map.get(map, Atom.to_string(key))
  end

  defp field(_map, _key), do: nil

  defp list_field(map, key) do
    case field(map, key) do
      list when is_list(list) -> list
      _ -> []
    end
  end

  defp numeric_field(map, key) do
    case field(map, key) do
      value when is_number(value) -> value
      _ -> nil
    end
  end

  defp maybe_assert_any(_items, _key, nil), do: :ok

  defp maybe_assert_any(items, key, expected) do
    assert Enum.any?(items, &(field(&1, key) == expected)),
           "expected any citation #{key}=#{inspect(expected)}, got #{inspect(items)}"
  end

  defp present?(value) when is_binary(value), do: String.trim(value) != ""
  defp present?(value), do: not is_nil(value)

  defp collect_raw_text_keys(%{} = map, path) do
    Enum.flat_map(map, fn {key, value} ->
      normalized = normalize_key(key)
      current_path = path ++ [normalized]

      if normalized in @raw_text_string_keys and raw_value?(value) do
        [Enum.join(current_path, ".")]
      else
        collect_raw_text_keys(value, current_path)
      end
    end)
  end

  defp collect_raw_text_keys(list, path) when is_list(list) do
    list
    |> Enum.with_index()
    |> Enum.flat_map(fn {value, index} ->
      collect_raw_text_keys(value, path ++ [to_string(index)])
    end)
  end

  defp collect_raw_text_keys(_value, _path), do: []

  defp raw_value?(nil), do: false
  defp raw_value?(""), do: false
  defp raw_value?([]), do: false
  defp raw_value?(%{} = map), do: map_size(map) > 0
  defp raw_value?(_value), do: true

  defp normalize_key(key) when is_atom(key), do: Atom.to_string(key)
  defp normalize_key(key), do: to_string(key)
end