# SPDX-FileCopyrightText: 2025 Isaak Tsalicoglou <isaak@overbring.com>
# SPDX-License-Identifier: BSD-3-Clause
defmodule DoTrace do
  @moduledoc """
  Provides a `trace/1` macro to trace function calls at runtime,
  mimicking Clojure's `dotrace` output.

  ## Example usage

  ```elixir
  iex> defmodule Math do
         def factorial(0), do: 1
         def factorial(n) when n > 0, do: n * factorial(n - 1)
       end
  iex> require DoTrace
  iex> DoTrace.trace(Math.factorial(5))
  TRACE t1028: (factorial [5])
  TRACE t1092: | (factorial [4])
  TRACE t1156: | | (factorial [3])
  TRACE t1220: | | | (factorial [2])
  TRACE t1284: | | | | (factorial [1])
  TRACE t1348: | | | | | (factorial [0])
  TRACE t1348: | | | | | | => 1
  TRACE t1284: | | | | | => 1
  TRACE t1220: | | | | => 2
  TRACE t1156: | | | => 6
  TRACE t1092: | | => 24
  TRACE t1028: | => 120
  120
  ```
  """

  @spec trace({atom() | {:., any(), [...]}, any(), any()}) ::
          {:__block__, [], [{:= | {any(), any(), any()}, [...], [...]}, ...]}
  defmacro trace(call_ast) do
    {m, f, args_ast} = parse_call(call_ast, __CALLER__)

    quote do
      args = [unquote_splicing(args_ast)]
      DoTrace.Tracer.run(unquote(m), unquote(f), args)
    end
  end

  defp parse_call({{:., _, [module, fun]}, _, args}, _caller) when is_atom(fun) do
    {module, fun, args}
  end

  defp parse_call({fun, _, args}, caller) when is_atom(fun) do
    {caller.module, fun, args}
  end

  defmodule Tracer do
    @moduledoc false
    @spec run(atom(), atom(), list()) :: any()
    def run(mod, fun, args) do
      arity = length(args)
      mfa = {mod, fun, arity}
      tracer = self()

      :erlang.trace_pattern({:_, :_, :_}, false, [:local])
      :erlang.trace_pattern(mfa, false, [:local])
      :erlang.trace(tracer, true, [:call, :set_on_spawn])
      :erlang.trace_pattern(mfa, [{:_, [], [{:return_trace}]}], [:local])

      task =
        Task.async(fn ->
          Process.sleep(10)
          apply(mod, fun, args)
        end)

      if Process.alive?(task.pid) do
        :erlang.trace(task.pid, true, [:call])
      else
        IO.puts(
          :stderr,
          "Warning: Task PID #{inspect(task.pid)} is not alive, skipping trace setup"
        )
      end

      monitor_ref = Process.monitor(task.pid)
      flush_mailbox()

      initial_state = %{depth: 0, calls: %{}, result: nil}
      final_result = loop_and_format_traces(task.ref, task.pid, mfa, initial_state)

      :erlang.trace_pattern(mfa, false, [:local])
      :erlang.trace(tracer, false, [:call])

      if Process.alive?(task.pid) do
        :erlang.trace(task.pid, false, [:call])
      end

      Process.demonitor(monitor_ref, [:flush])
      flush_mailbox()

      final_result
    end

    defp flush_mailbox do
      receive do
        _msg ->
          # IO.inspect(msg, label: "Flushed Message")
          flush_mailbox()
      after
        0 -> :ok
      end
    end

    defp loop_and_format_traces(task_ref, task_pid, mfa, state) do
      receive do
        {:trace, pid, :call, {m, f, args}} when pid == task_pid and {m, f, length(args)} == mfa ->
          trace_id = :erlang.unique_integer([:positive])
          indent = String.duplicate("| ", state.depth)

          IO.puts(
            :stdio,
            "TRACE t#{trace_id}: #{indent}(#{f} #{inspect(args, charlists: :as_lists)})"
          )

          new_state =
            state
            |> update_in([:depth], &(&1 + 1))
            |> put_in([:calls, state.depth + 1], {trace_id, nil})

          loop_and_format_traces(task_ref, task_pid, mfa, new_state)

        {:trace, pid, :return_from, {m, f, a}, result}
        when pid == task_pid and {m, f, a} == mfa ->
          new_depth = max(state.depth - 1, 0)
          {trace_id, _} = get_in(state, [:calls, state.depth]) || {0, nil}
          indent = String.duplicate("| ", state.depth)

          IO.puts(
            :stdio,
            "TRACE t#{trace_id}: #{indent}=> #{inspect(result, charlists: :as_lists)}"
          )

          new_state =
            state
            |> Map.put(:depth, new_depth)
            |> Map.put(:calls, Map.delete(state.calls, state.depth))

          loop_and_format_traces(task_ref, task_pid, mfa, new_state)

        {^task_ref, result} ->
          new_state = Map.put(state, :result, result)
          loop_and_format_traces(task_ref, task_pid, mfa, new_state)

        {:DOWN, _ref, :process, ^task_pid, _status} ->
          state.result

        _other ->
          # IO.inspect(other, label: "Unexpected Message")
          loop_and_format_traces(task_ref, task_pid, mfa, state)
      end
    end
  end
end
