cmdc_test

defmodule CMDCTest.Assertions do
  @moduledoc """
  CMDC EventBus 事件断言族 — 替代手写 receive/after 模板。

  必须先调 `CMDCTest.EventCapture.start_capture/1` 让当前进程订阅事件。

  ## API

  - `assert_event_emitted/3` — 期望收到指定事件
  - `refute_event_emitted/3` — 期望**不**收到指定事件
  - `assert_event_count/4` — 期望恰好收到 N 个指定事件

  匹配语义：事件首元 atom 等于目标 atom 即匹配；payload 可选传 map 做子集匹配。

  ## 不匹配事件的处理

  函数内部 receive 到的**不匹配事件会重新 push 回当前进程 mailbox**，
  确保后续断言能继续看到它们。顺序会按 push 时序重排到 mailbox 尾部，
  对绝大多数测试断言无影响。

  ## Quick Start

      import CMDCTest.Assertions

      assert_event_emitted(session, :agent_end)
      assert_event_emitted(session, :tool_blocked, payload: %{tool: "shell"})
      assert_event_emitted(session, :token, timeout: 1_000)

      refute_event_emitted(session, :approval_required, timeout: 200)

      events = assert_event_count(session, :stream_chunk, 5, timeout: 1_000)
  """

  alias CMDCTest.EventCapture

  import ExUnit.Assertions

  @default_assert_timeout 500
  @default_refute_timeout 200

  # ==========================================================================
  # assert_event_emitted/2,3
  # ==========================================================================

  @doc """
  断言当前进程在 timeout 内收到指定事件，返回匹配的 event term。

  ## 选项

  - `:timeout` — 最大等待时间（默认 500ms）
  - `:payload` — payload map 子集匹配（默认 `:_` 不校验）
  """
  @spec assert_event_emitted(CMDC.session(), atom(), keyword()) :: term()
  def assert_event_emitted(session, event_atom, opts \\ []) when is_atom(event_atom) do
    {:ok, sid} = EventCapture.normalize_sid(session)
    timeout = Keyword.get(opts, :timeout, @default_assert_timeout)
    payload_match = Keyword.get(opts, :payload, :_)

    case wait_for(sid, event_atom, payload_match, timeout, []) do
      {:ok, event} ->
        event

      :timeout ->
        flunk(
          "expected event #{inspect(event_atom)} not received within #{timeout}ms " <>
            "(session=#{sid}). Did you forget to call CMDCTest.EventCapture.start_capture/1?"
        )
    end
  end

  # ==========================================================================
  # refute_event_emitted/2,3
  # ==========================================================================

  @doc """
  断言当前进程在 timeout 内**不**收到指定事件。

  ## 选项

  - `:timeout` — 等待时间（默认 200ms）
  - `:payload` — payload map 子集匹配
  """
  @spec refute_event_emitted(CMDC.session(), atom(), keyword()) :: :ok
  def refute_event_emitted(session, event_atom, opts \\ []) when is_atom(event_atom) do
    {:ok, sid} = EventCapture.normalize_sid(session)
    timeout = Keyword.get(opts, :timeout, @default_refute_timeout)
    payload_match = Keyword.get(opts, :payload, :_)

    case wait_for(sid, event_atom, payload_match, timeout, []) do
      {:ok, event} ->
        flunk("expected NOT to receive event #{inspect(event_atom)}, but got #{inspect(event)}")

      :timeout ->
        :ok
    end
  end

  # ==========================================================================
  # assert_event_count/3,4
  # ==========================================================================

  @doc """
  断言恰好收到 N 个指定事件（在 timeout 内）。

  返回收到的事件列表（按时序排列）。
  """
  @spec assert_event_count(CMDC.session(), atom(), non_neg_integer(), keyword()) :: [term()]
  def assert_event_count(session, event_atom, n, opts \\ [])
      when is_atom(event_atom) and is_integer(n) and n >= 0 do
    {:ok, sid} = EventCapture.normalize_sid(session)
    timeout = Keyword.get(opts, :timeout, @default_assert_timeout)
    payload_match = Keyword.get(opts, :payload, :_)

    events = collect_n(sid, event_atom, payload_match, n, timeout, [])
    actual = length(events)

    assert actual == n,
           "expected #{n} #{inspect(event_atom)} events within #{timeout}ms, got #{actual}"

    events
  end

  # ==========================================================================
  # 私有：receive loop + 不匹配事件重排
  # ==========================================================================

  defp wait_for(sid, target, pm, timeout, drained) do
    deadline = System.monotonic_time(:millisecond) + timeout

    case receive_one(sid, target, pm, deadline, drained) do
      {:found, event, leftover} ->
        requeue(leftover)
        {:ok, event}

      {:timeout, leftover} ->
        requeue(leftover)
        :timeout
    end
  end

  defp receive_one(sid, target, pm, deadline, drained) do
    remaining = max(deadline - System.monotonic_time(:millisecond), 0)

    if remaining == 0 do
      {:timeout, drained}
    else
      receive do
        {:cmdc_event, ^sid, event} = msg ->
          if matches?(event, target, pm) do
            {:found, event, drained}
          else
            receive_one(sid, target, pm, deadline, [msg | drained])
          end
      after
        remaining -> {:timeout, drained}
      end
    end
  end

  defp collect_n(_sid, _target, _pm, n, _timeout, acc) when length(acc) >= n,
    do: Enum.reverse(acc)

  defp collect_n(sid, target, pm, n, timeout, acc) do
    case wait_for(sid, target, pm, timeout, []) do
      {:ok, event} -> collect_n(sid, target, pm, n, timeout, [event | acc])
      :timeout -> Enum.reverse(acc)
    end
  end

  defp requeue(drained) do
    drained |> Enum.reverse() |> Enum.each(&send(self(), &1))
  end

  # ==========================================================================
  # 匹配规则
  # ==========================================================================

  # bare atom event（无 payload）
  defp matches?(target, target, :_) when is_atom(target), do: true

  # tuple-form: 首元 atom == target
  defp matches?(event_tuple, target, :_)
       when is_tuple(event_tuple) and tuple_size(event_tuple) >= 1 do
    elem(event_tuple, 0) == target
  end

  # tuple-form with payload match
  defp matches?(event_tuple, target, expected_payload)
       when is_tuple(event_tuple) and tuple_size(event_tuple) >= 2 do
    elem(event_tuple, 0) == target and
      payload_subset?(extract_payload(event_tuple), expected_payload)
  end

  defp matches?(_, _, _), do: false

  # 多元 tuple event 的 payload 通常在第二元素
  defp extract_payload(event_tuple), do: elem(event_tuple, 1)

  defp payload_subset?(actual, expected) when is_map(actual) and is_map(expected) do
    Enum.all?(expected, fn {k, v} -> Map.get(actual, k) == v end)
  end

  defp payload_subset?(actual, expected), do: actual == expected
end