defmodule Qtrace do
  @moduledoc """
  Quick function latency analysis and visualization based on Erlang tracing.

  `Qtrace` traces function calls and collects latency measurements. Uses
  histogram-based aggregation to provide quantile analysis (P50, P95, P99)
  and terminal-based visualization of performance distributions.

  ## Usage

      session = Qtrace.start_session(:all)
      Qtrace.trace_function(session, MyApp, :my_function, 2)
      Qtrace.trace_function(session, MyApp, :another_function, 1)

      # ... run your code ...

      Qtrace.plot_all_histograms(session)

      # Getting p10
      Qtrace.get_quantile(session, 0.10, MyApp, :my_function, 2)

  ## Example output

  ```txt
  Latency Histogram: Qtrace.Demo.web_request/0
  ================================================================================
              0              13             26             39             52
        5.9ms  │███ (3)
        6.9ms  │█ (1)
        8.1ms  │██ (2)
        10.3ms │█ (1)
        11.2ms │██████ (5)
        12.1ms │█████ (4)
        13.2ms │███ (3)
        14.2ms │███████ (6)
        15.4ms │████████████████ (14)
        16.7ms │█████████████████ (15)
        18.1ms │████████████████████ (17)
        19.6ms │█████████████████████████████████ (29)
        21.3ms │██████████████████████████████████████████████████████████ (50)
        23.0ms │██████████████████████████████████████████ (36)
        24.9ms │███████████████████████████████████████████████ (41)
        27.0ms │████████████████████████████████████████████████████████████ (52)
        29.3ms │█████████████████████████████████████████████████████████ (49)
        31.7ms │███████████████████████████████████████████████████████████ (51)
        34.4ms │██████████████████████████████████████ (33)
        37.2ms │████████████████████████████████████████ (35)
        40.3ms │████████████████████████ (21)
        43.7ms │███████ (6)
        47.3ms │███ (3)
        51.3ms │██ (2)
      134.0ms  │██ (2)
      157.2ms  │█ (1)
      199.9ms  │█ (1)
      216.6ms  │█████████ (8)
      234.6ms  │██ (2)
      254.2ms  │██ (2)
      275.4ms  │███ (3)
      298.3ms  │██ (2)
      323.2ms  │██ (2)
               └────────────────────────────────────────────────────────────
              0              13             26             39             52
  Stats:
    Min: 5.9ms
    Max: 323.2ms
    Samples: 502
    P50: 27.0ms
    P95: 47.3ms
    P99: 275.4ms
  ```

  ## Session Management

  Tracing is automatically stopped when the trace session gets garbage collected.
  The trace process may live for a bit longer: it lives for up to two minutes of
  inactivity.

  You can also stop sessions or clean up trace data manually:

      Qtrace.stop_session(session) # Stop the trace session
      Qtrace.cleanup(session)      # Stop the tracer process that stores the results

  Notice that, in some situations, **you may need to ensure that the session doesn't
  get garbage collected** to avoid the trace stopping unexpectedly.

  ## Notes

  - Sessions auto-stop when querying results
  - Measurements are in microseconds
  - Buckets have an error of up to 4%
  """

  alias Qtrace.Histogram
  alias Qtrace.Tracer
  alias Qtrace.Session

  @untraceable_modules [
    :tracer,
    Qtrace,
    Qtrace.Session,
    Qtrace.Tracer,
    DogSketch.SimpleDog
  ]

  require Logger

  @doc """
  Starts a new tracing session.

  ## Parameters

  - `processes` - Which processes to trace (`:all`, `:existing`, `:new`, or a `pid()`)

  ## Returns

  Returns a `Qtrace.Session.t()` struct containing the session details. If this struct
  is garbage collected, the trace stops.

  ## Examples

      # Trace all processes
      session = Qtrace.start_session(:all)

      # Trace only existing processes with a custom name
      session = Qtrace.start_session(:existing, :my_trace_session)

      # Trace a specific process
      session = Qtrace.start_session(self())

  """
  @spec start_session(:all | :existing | :new | pid()) :: Session.t()
  def start_session(processes) do
    :ok = validate_processes!(processes)

    {:ok, tracer} = Tracer.start_supervised([])

    ts = :trace.session_create(Qtrace.DefaultSession, tracer, [])

    :trace.process(ts, processes, true, [:call, :monotonic_timestamp])

    %Session{name: Qtrace.DefaultSession, trace_session: ts, pid: tracer}
  end

  @doc """
  Adds processes to an existing tracing session.

  ## Example

      session = Qtrace.start_session(:existing)
      Qtrace.add_processes(session, :new)  # Now traces existing + new processes

  """
  @spec add_processes(Session.t(), :all | :existing | :new | pid()) :: non_neg_integer()
  def add_processes(%Session{trace_session: ts}, processes) do
    :ok = validate_processes!(processes)
    :trace.process(ts, processes, true, [:call, :monotonic_timestamp])
  end

  @doc """
  Removes processes from an existing tracing session.

  ## Example

      session = Qtrace.start_session(:all)
      Qtrace.remove_processes(session, :new)  # Stop tracing new processes

  """
  @spec remove_processes(Session.t(), :all | :existing | :new | pid()) :: non_neg_integer()
  def remove_processes(%Session{trace_session: ts}, processes) do
    :ok = validate_processes!(processes)
    :trace.process(ts, processes, false, [:call, :monotonic_timestamp])
  end

  @doc """
  Enables tracing for a specific function.

  ## Example

      session = Qtrace.start_session(:all)
      Qtrace.trace_function(session, Enum, :map, 2)

      # Now all calls to Enum.map/2 will be traced

  ## Restrictions

  Some modules cannot be traced:
  - `:tracer`
  - `Qtrace` and its submodules
  - `DogSketch.SimpleDog`
  """
  @spec trace_function(Session.t(), module(), atom(), non_neg_integer()) :: non_neg_integer()
  def trace_function(%Session{trace_session: ts}, module, function, arity) do
    Code.ensure_loaded!(module)

    if module in @untraceable_modules do
      raise ArgumentError, """
      module #{inspect(module)} can't be traced by Qtrace

      Qtrace can't trace functions that it calls itself, otherwise we could get a trace loop\
      """
    end

    :trace.function(
      ts,
      {module, function, arity},
      [{:_, [], [{:return_trace}]}],
      [:local]
    )
  end

  @doc """
  Gets a quantile from traced function performance data. Returns latency in microseconds.

  ## Examples

      # Get median latency (P50)
      p50 = Qtrace.get_quantile(session, 0.5, MyApp, :my_function, 1)

      # Get 95th percentile
      p95 = Qtrace.get_quantile(session, 0.95, MyApp, :my_function, 1)
  """
  @spec get_quantile(
          Session.t(),
          quantile :: float(),
          module(),
          function_name :: atom(),
          arity()
        ) ::
          number() | nil
  def get_quantile(session, quantile, module, function, arity) do
    auto_destroy_session(session)

    if quantile <= 0 or quantile >= 1 do
      raise ArgumentError, "quantile must be between 0 and 1, got #{quantile}"
    end

    Tracer.get_quantile(session.pid, quantile, module, function, arity)
  end

  @doc """
  Builds histogram data from traced function measurements.

  Returns a `%Qtrace.Histogram{}` struct or `nil` if no data available.
  """
  @spec get_histogram(Session.t(), module(), atom(), non_neg_integer()) ::
          Qtrace.Histogram.t() | nil
  def get_histogram(session, module, function, arity) do
    auto_destroy_session(session)

    case Tracer.get_sketch(session.pid, module, function, arity) do
      nil -> nil
      sketch -> Histogram.build_histogram_from_sketch(sketch, module, function, arity)
    end
  end

  @doc """
  Displays a histogram of latency measurements in the terminal.

  ### Options

  - `:width` - Width of the histogram bars (default: 60)

  ## Example

      Qtrace.plot_histogram(session, MyApp, :my_function, 1)

  """
  @spec plot_histogram(Session.t(), module(), atom(), non_neg_integer(), keyword()) :: :ok
  def plot_histogram(session, module, function, arity, opts \\ []) do
    auto_destroy_session(session)

    case get_histogram(session, module, function, arity) do
      nil ->
        IO.puts("No data available for #{inspect(module)}.#{function}/#{arity}")

      histogram ->
        Qtrace.Plotter.plot_histogram(histogram, opts)
    end
  end

  @doc """
  Displays histograms for all traced functions.

  ### Options

  - `:width` - Width of the histogram bars (default: 60)
  """
  @spec plot_all_histograms(Session.t(), keyword()) :: :ok
  def plot_all_histograms(session, opts \\ []) do
    auto_destroy_session(session)

    case Tracer.list_functions(session.pid) do
      [] ->
        IO.puts("No traced functions found")
        :ok

      functions ->
        Enum.each(functions, fn {module, function, arity} ->
          plot_histogram(session, module, function, arity, opts)
        end)
    end
  end

  @doc """
  Stops a tracing session. Trace results remain available until cleanup.

  ## Example

      session = Qtrace.start_session(:all)
      # ... do tracing work ...
      Qtrace.stop_session(session)

  """
  @spec stop_session(Session.t()) :: Session.t()
  def stop_session(%Session{} = session) do
    :trace.session_destroy(session.trace_session)
    session
  end

  @doc """
  Stops the tracer process for a session.

  ## Example

      session = Qtrace.start_session(:all)
      # ... do tracing work ...
      Qtrace.cleanup(session)  # Immediately stop the tracer

  """
  @spec cleanup(Session.t()) :: :ok
  def cleanup(%Session{} = session) do
    Tracer.stop(session.pid)
  end

  defp validate_processes!(:all), do: :ok
  defp validate_processes!(:existing), do: :ok
  defp validate_processes!(:new), do: :ok

  defp validate_processes!(pid) when is_pid(pid) do
    if node(pid) == Node.self() do
      :ok
    else
      raise "invalid process: #{inspect(pid)}, should be a local process"
    end
  end

  defp validate_processes!(other) do
    raise ArgumentError,
          "invalid processes #{inspect(other)}, should be either: :all, :existing, :new or a local pid"
  end

  defp auto_destroy_session(session) do
    if :trace.session_destroy(session.trace_session) do
      Logger.warning(
        "Stopped trace session #{inspect(session.name)}: can't query results with session active"
      )
    end
  end
end