firebird

defmodule Firebird.Profiler do
  @moduledoc """
  Profiling tools for WASM execution performance analysis.

  Provides detailed timing information for WASM function calls,
  including startup time, call overhead, and throughput metrics.

  ## Examples

      # Profile a single function
      profile = Firebird.Profiler.profile("math.wasm", :fibonacci, [30], iterations: 1000)
      IO.puts("Avg: \#{profile.avg_us}μs, p99: \#{profile.p99_us}μs")

      # Profile module startup
      startup = Firebird.Profiler.profile_startup("math.wasm", iterations: 50)
      IO.puts("Load time: \#{startup.avg_us}μs")

      # Compare multiple implementations
      comparison = Firebird.Profiler.compare_implementations(%{
        rust: {"fixtures/rust_math.wasm", []},
        go: {"fixtures/go_math.wasm", [wasi: true]},
        native: fn args -> apply(&Kernel.+/2, args) end
      }, :add, [5, 3])

      # Full module profile
      report = Firebird.Profiler.profile_module("math.wasm")
  """

  @default_iterations 100
  @default_warmup 10

  @doc """
  Profile a specific WASM function call.

  ## Options

  - `:iterations` - Number of measured iterations (default: 100)
  - `:warmup` - Number of warmup iterations (default: 10)
  - `:wasi` - Enable WASI (default: false)

  ## Returns

  A map with timing statistics in microseconds.
  """
  @spec profile(String.t(), atom(), list(), keyword()) :: map()
  def profile(wasm_path, function, args, opts \\ []) do
    iterations = Keyword.get(opts, :iterations, @default_iterations)
    warmup = Keyword.get(opts, :warmup, @default_warmup)
    load_opts = Keyword.drop(opts, [:iterations, :warmup])

    {:ok, instance} = Firebird.WasmRunner.start(wasm_path, load_opts)

    # Warmup
    for _ <- 1..warmup do
      Firebird.call(instance, function, args)
    end

    # Measure
    times =
      for _ <- 1..iterations do
        {time, _} = :timer.tc(fn -> Firebird.call(instance, function, args) end)
        time
      end

    Firebird.stop(instance)
    compute_stats(times, to_string(function), iterations)
  end

  @doc """
  Profile WASM module startup (load) time.

  ## Options

  - `:iterations` - Number of load/unload cycles (default: 20)
  - `:wasi` - Enable WASI (default: false)
  """
  @spec profile_startup(String.t(), keyword()) :: map()
  def profile_startup(wasm_path, opts \\ []) do
    iterations = Keyword.get(opts, :iterations, 20)
    load_opts = Keyword.drop(opts, [:iterations])

    times =
      for _ <- 1..iterations do
        {time, {:ok, instance}} =
          :timer.tc(fn ->
            Firebird.WasmRunner.start(wasm_path, load_opts)
          end)

        Firebird.stop(instance)
        time
      end

    stats = compute_stats(times, "startup", iterations)
    Map.put(stats, :wasm_path, wasm_path)
  end

  @doc """
  Profile all exported functions in a WASM module.

  Automatically discovers exports and profiles each one with
  default arguments (zeros).

  ## Options

  - `:iterations` - Iterations per function (default: 50)
  - `:wasi` - Enable WASI (default: false)
  """
  @spec profile_module(String.t(), keyword()) :: map()
  def profile_module(wasm_path, opts \\ []) do
    iterations = Keyword.get(opts, :iterations, 50)
    load_opts = Keyword.drop(opts, [:iterations])

    {:ok, instance} = Firebird.WasmRunner.start(wasm_path, load_opts)

    sigs = Firebird.WasmRunner.function_signatures(instance)

    profiles =
      Enum.map(sigs, fn {func, {params, _results}} ->
        # Generate default args (all zeros)
        args = Enum.map(params, fn _ -> 0 end)

        # Skip functions that might trap with zero args
        times =
          try do
            # Warmup
            for _ <- 1..5, do: Firebird.call(instance, func, args)

            for _ <- 1..iterations do
              {time, _} = :timer.tc(fn -> Firebird.call(instance, func, args) end)
              time
            end
          rescue
            _ -> []
          catch
            _, _ -> []
          end

        if times != [] do
          {func, compute_stats(times, to_string(func), iterations)}
        else
          {func, %{function: to_string(func), error: "call failed with default args"}}
        end
      end)

    Firebird.stop(instance)

    %{
      wasm_path: wasm_path,
      function_count: map_size(sigs),
      profiles: Map.new(profiles),
      summary: summarize_profiles(profiles)
    }
  end

  @doc """
  Compare performance of the same function across different WASM implementations
  and optionally a native Elixir function.

  ## Examples

      result = Firebird.Profiler.compare_implementations(%{
        rust: {"fixtures/rust_math.wasm", []},
        go: {"fixtures/go_math.wasm", [wasi: true]},
        native: fn [a, b] -> a + b end
      }, :add, [5, 3], iterations: 200)
  """
  @spec compare_implementations(map(), atom(), list(), keyword()) :: map()
  def compare_implementations(implementations, function, args, opts \\ []) do
    iterations = Keyword.get(opts, :iterations, @default_iterations)
    warmup = Keyword.get(opts, :warmup, @default_warmup)

    results =
      Enum.map(implementations, fn
        {name, {wasm_path, load_opts}} ->
          {:ok, instance} = Firebird.WasmRunner.start(wasm_path, load_opts)

          # Warmup
          for _ <- 1..warmup, do: Firebird.call(instance, function, args)

          times =
            for _ <- 1..iterations do
              {time, _} = :timer.tc(fn -> Firebird.call(instance, function, args) end)
              time
            end

          Firebird.stop(instance)
          {name, compute_stats(times, "#{name}/#{function}", iterations)}

        {name, func} when is_function(func) ->
          # Native Elixir function
          for _ <- 1..warmup, do: func.(args)

          times =
            for _ <- 1..iterations do
              {time, _} = :timer.tc(fn -> func.(args) end)
              time
            end

          {name, compute_stats(times, "#{name}/#{function}", iterations)}
      end)

    results_map = Map.new(results)
    fastest = results |> Enum.min_by(fn {_, stats} -> stats.avg_us end) |> elem(0)

    %{
      function: function,
      args: args,
      iterations: iterations,
      results: results_map,
      fastest: fastest,
      rankings: rank_results(results)
    }
  end

  @doc """
  Measure throughput: how many calls per second.

  ## Examples

      throughput = Firebird.Profiler.throughput("math.wasm", :add, [5, 3], duration_ms: 1000)
      IO.puts("\#{throughput.calls_per_second} calls/sec")
  """
  @spec throughput(String.t(), atom(), list(), keyword()) :: map()
  def throughput(wasm_path, function, args, opts \\ []) do
    duration_ms = Keyword.get(opts, :duration_ms, 1000)
    load_opts = Keyword.drop(opts, [:duration_ms])

    {:ok, instance} = Firebird.WasmRunner.start(wasm_path, load_opts)

    # Warmup
    for _ <- 1..10, do: Firebird.call(instance, function, args)

    # Run for duration
    deadline = System.monotonic_time(:millisecond) + duration_ms
    count = count_calls(instance, function, args, deadline, 0)

    Firebird.stop(instance)

    %{
      function: function,
      duration_ms: duration_ms,
      total_calls: count,
      calls_per_second: count * 1000 / duration_ms,
      avg_us: duration_ms * 1000 / max(count, 1)
    }
  end

  defp count_calls(instance, function, args, deadline, count) do
    if System.monotonic_time(:millisecond) >= deadline do
      count
    else
      Firebird.call(instance, function, args)
      count_calls(instance, function, args, deadline, count + 1)
    end
  end

  @doc """
  Format profiler results as a readable string.
  """
  @spec format(map()) :: String.t()
  def format(%{results: results, fastest: fastest} = comparison) do
    lines = [
      "📊 Performance Comparison: #{comparison.function}(#{Enum.join(comparison.args, ", ")})",
      "   Iterations: #{comparison.iterations}",
      ""
    ]

    result_lines =
      Enum.map(results, fn {name, stats} ->
        marker = if name == fastest, do: "🏆", else: "  "

        "   #{marker} #{String.pad_trailing("#{name}", 15)} avg=#{Float.round(stats.avg_us, 1)}μs  p50=#{Float.round(stats.p50_us, 1)}μs  p99=#{Float.round(stats.p99_us, 1)}μs"
      end)

    Enum.join(lines ++ result_lines, "\n")
  end

  def format(%{profiles: profiles, summary: summary} = module_profile) do
    lines = [
      "📊 Module Profile: #{module_profile.wasm_path}",
      "   Functions: #{module_profile.function_count}",
      "   Avg call time: #{Float.round(summary.avg_call_us, 1)}μs",
      "   Fastest: #{summary.fastest} (#{Float.round(summary.fastest_us, 1)}μs)",
      "   Slowest: #{summary.slowest} (#{Float.round(summary.slowest_us, 1)}μs)",
      ""
    ]

    func_lines =
      profiles
      |> Enum.reject(fn {_, stats} -> Map.has_key?(stats, :error) end)
      |> Enum.sort_by(fn {_, stats} -> stats.avg_us end)
      |> Enum.map(fn {name, stats} ->
        "   #{String.pad_trailing("#{name}", 20)} avg=#{Float.round(stats.avg_us, 1)}μs"
      end)

    Enum.join(lines ++ func_lines, "\n")
  end

  # --- Private ---

  defp compute_stats(times, name, _iterations) do
    sorted = Enum.sort(times)
    count = length(sorted)

    if count == 0 do
      %{
        function: name,
        avg_us: 0.0,
        min_us: 0,
        max_us: 0,
        p50_us: 0.0,
        p95_us: 0.0,
        p99_us: 0.0,
        stddev_us: 0.0
      }
    else
      sum = Enum.sum(sorted)
      avg = sum / count

      %{
        function: name,
        avg_us: avg,
        min_us: hd(sorted),
        max_us: List.last(sorted),
        median_us: percentile(sorted, 50),
        p50_us: percentile(sorted, 50),
        p95_us: percentile(sorted, 95),
        p99_us: percentile(sorted, 99),
        stddev_us: stddev(sorted, avg),
        total_us: sum,
        count: count
      }
    end
  end

  defp percentile(sorted, p) when length(sorted) > 0 do
    k = p / 100 * (length(sorted) - 1)
    f = trunc(k)
    c = min(f + 1, length(sorted) - 1)
    lower = Enum.at(sorted, f)
    upper = Enum.at(sorted, c)
    lower + (upper - lower) * (k - f)
  end

  defp percentile(_, _), do: 0.0

  defp stddev(values, mean) when length(values) > 1 do
    variance =
      Enum.reduce(values, 0.0, fn v, acc ->
        diff = v - mean
        acc + diff * diff
      end) / (length(values) - 1)

    :math.sqrt(variance)
  end

  defp stddev(_, _), do: 0.0

  defp summarize_profiles(profiles) do
    valid =
      profiles
      |> Enum.reject(fn {_, stats} -> Map.has_key?(stats, :error) end)
      |> Enum.map(fn {name, stats} -> {name, stats.avg_us} end)

    if valid == [] do
      %{avg_call_us: 0.0, fastest: nil, fastest_us: 0.0, slowest: nil, slowest_us: 0.0}
    else
      avg = valid |> Enum.map(&elem(&1, 1)) |> Enum.sum() |> Kernel./(length(valid))
      {fastest, fastest_us} = Enum.min_by(valid, &elem(&1, 1))
      {slowest, slowest_us} = Enum.max_by(valid, &elem(&1, 1))

      %{
        avg_call_us: avg,
        fastest: fastest,
        fastest_us: fastest_us,
        slowest: slowest,
        slowest_us: slowest_us
      }
    end
  end

  defp rank_results(results) do
    results
    |> Enum.sort_by(fn {_, stats} -> stats.avg_us end)
    |> Enum.with_index(1)
    |> Enum.map(fn {{name, stats}, rank} -> {name, rank, stats.avg_us} end)
  end
end