firebird

defmodule Firebird.Target.Profiler do
  @moduledoc """
  Profile the WASM compilation pipeline to identify bottlenecks.

  Measures the time spent in each pipeline stage and produces
  a detailed timing report.

  ## Usage

      {:ok, profile} = Profiler.profile_source(source_code)
      IO.puts(Profiler.format(profile))

      # Profile a file
      {:ok, profile} = Profiler.profile_file("lib/math.ex")

      # Profile with specific options
      {:ok, profile} = Profiler.profile_source(source, optimize: true, tco: true)
  """

  @stages [
    :read,
    :parse,
    :ir_gen,
    :optimize,
    :tco,
    :inline,
    :validate,
    :type_infer,
    :wat_gen,
    :wat2wasm
  ]

  @type timing :: %{
          stage: atom(),
          elapsed_us: non_neg_integer(),
          percentage: float()
        }

  @type profile :: %{
          source: String.t() | nil,
          file: String.t() | nil,
          module: atom() | nil,
          timings: [timing()],
          total_us: non_neg_integer(),
          result: :ok | :error,
          error: term() | nil,
          options: keyword()
        }

  @doc """
  Profile compilation of a source code string.

  Returns detailed timing for each pipeline stage.

  ## Options

  - `:optimize` - Enable optimization pass (default: false)
  - `:tco` - Enable TCO pass (default: false)
  - `:inline` - Enable inlining pass (default: false)
  - `:wat_only` - Skip binary compilation (default: false)
  """
  @spec profile_source(String.t(), keyword()) :: {:ok, profile()}
  def profile_source(source, opts \\ []) do
    optimize = Keyword.get(opts, :optimize, false)
    tco = Keyword.get(opts, :tco, false)
    do_inline = Keyword.get(opts, :inline, false)
    wat_only = Keyword.get(opts, :wat_only, false)

    timings = []

    # Parse
    {parse_us, parse_result} =
      time_stage(fn ->
        Firebird.Compiler.parse(source)
      end)

    timings = [{:parse, parse_us} | timings]

    case parse_result do
      {:ok, ast} ->
        # IR Generation
        {ir_us, ir_result} =
          time_stage(fn ->
            Firebird.Compiler.IRGen.generate(ast)
          end)

        timings = [{:ir_gen, ir_us} | timings]

        case ir_result do
          {:ok, module_ir} ->
            # Optional: Optimize
            {opt_us, opt_result} =
              time_stage(fn ->
                if optimize do
                  Firebird.Compiler.Optimizer.optimize(module_ir)
                else
                  {:ok, module_ir}
                end
              end)

            timings = [{:optimize, opt_us} | timings]

            case opt_result do
              {:ok, opt_ir} ->
                # Optional: TCO
                {tco_us, tco_result} =
                  time_stage(fn ->
                    if tco do
                      Firebird.Compiler.TCO.optimize(opt_ir)
                    else
                      {:ok, opt_ir}
                    end
                  end)

                timings = [{:tco, tco_us} | timings]

                case tco_result do
                  {:ok, tco_ir} ->
                    # Optional: Inline
                    {inline_us, inline_result} =
                      time_stage(fn ->
                        if do_inline do
                          Firebird.Compiler.Inliner.inline(tco_ir)
                        else
                          {:ok, tco_ir}
                        end
                      end)

                    timings = [{:inline, inline_us} | timings]

                    case inline_result do
                      {:ok, final_ir} ->
                        # Validate
                        {val_us, val_result} =
                          time_stage(fn ->
                            Firebird.Compiler.Validator.validate(final_ir)
                          end)

                        timings = [{:validate, val_us} | timings]

                        case val_result do
                          :ok ->
                            # Type Inference
                            {type_us, type_result} =
                              time_stage(fn ->
                                Firebird.Compiler.TypeInference.infer(final_ir)
                              end)

                            timings = [{:type_infer, type_us} | timings]

                            case type_result do
                              {:ok, typed_ir} ->
                                # WAT Generation
                                {wat_us, wat_result} =
                                  time_stage(fn ->
                                    Firebird.Compiler.WATGen.generate(typed_ir)
                                  end)

                                timings = [{:wat_gen, wat_us} | timings]

                                case wat_result do
                                  {:ok, wat} ->
                                    # WAT to WASM
                                    {wasm_us, _wasm_result} =
                                      if wat_only do
                                        {0, {:ok, nil}}
                                      else
                                        time_stage(fn -> Firebird.Compiler.wat_to_wasm(wat) end)
                                      end

                                    timings = [{:wat2wasm, wasm_us} | timings]

                                    build_profile(timings, typed_ir.name, nil, opts, :ok)

                                  {:error, reason} ->
                                    build_profile(timings, nil, nil, opts, :error, reason)
                                end

                              {:error, reason} ->
                                build_profile(timings, nil, nil, opts, :error, reason)
                            end

                          {:error, reason} ->
                            build_profile(timings, nil, nil, opts, :error, reason)
                        end

                      {:error, reason} ->
                        build_profile(timings, nil, nil, opts, :error, reason)
                    end

                  {:error, reason} ->
                    build_profile(timings, nil, nil, opts, :error, reason)
                end

              {:error, reason} ->
                build_profile(timings, nil, nil, opts, :error, reason)
            end

          {:error, reason} ->
            build_profile(timings, nil, nil, opts, :error, reason)
        end

      {:error, reason} ->
        build_profile(timings, nil, nil, opts, :error, reason)
    end
  end

  @doc """
  Profile compilation of a file.
  """
  @spec profile_file(String.t(), keyword()) :: {:ok, profile()}
  def profile_file(path, opts \\ []) do
    {read_us, read_result} = time_stage(fn -> File.read(path) end)

    case read_result do
      {:ok, source} ->
        {:ok, profile} = profile_source(source, opts)
        # Convert existing timing maps back to tuples, prepend read, and rebuild
        existing_tuples = Enum.map(profile.timings, fn t -> {t.stage, t.elapsed_us} end)
        all_tuples = [{:read, read_us} | existing_tuples]
        total = Enum.map(all_tuples, fn {_, us} -> us end) |> Enum.sum()

        {:ok, %{profile | file: path, timings: build_timings(all_tuples, total), total_us: total}}

      {:error, reason} ->
        build_profile([{:read, read_us}], nil, path, opts, :error, {:file_error, reason})
    end
  end

  @doc """
  Format a profile as a human-readable string.
  """
  @spec format(profile()) :: String.t()
  def format(profile) do
    lines = ["⏱ Compilation Profile"]

    lines =
      if profile.file do
        lines ++ ["  File: #{profile.file}"]
      else
        lines
      end

    lines =
      if profile.module do
        lines ++ ["  Module: #{profile.module}"]
      else
        lines
      end

    lines = lines ++ ["  Status: #{if profile.result == :ok, do: "✅ Success", else: "❌ Error"}"]
    lines = lines ++ [""]

    stage_lines =
      Enum.map(profile.timings, fn t ->
        name = t.stage |> Atom.to_string() |> String.pad_trailing(12)
        time = format_time(t.elapsed_us) |> String.pad_leading(10)
        pct = "#{Float.round(t.percentage, 1)}%" |> String.pad_leading(6)
        bar = String.duplicate("█", round(t.percentage / 2))
        "  #{name} #{time} #{pct} #{bar}"
      end)

    lines = lines ++ stage_lines
    lines = lines ++ ["", "  Total: #{format_time(profile.total_us)}"]

    if profile.error do
      lines = lines ++ ["  Error: #{inspect(profile.error)}"]
      Enum.join(lines, "\n")
    else
      Enum.join(lines, "\n")
    end
  end

  @doc """
  Return the list of all pipeline stages.
  """
  @spec stages() :: [atom()]
  def stages, do: @stages

  # Private helpers

  defp time_stage(fun) do
    :timer.tc(fun)
  end

  defp build_profile(timings, module, file, opts, result, error \\ nil) do
    timings = Enum.reverse(timings)
    total = Enum.map(timings, fn {_, us} -> us end) |> Enum.sum()

    {:ok,
     %{
       source: nil,
       file: file,
       module: module,
       timings: build_timings(timings, total),
       total_us: total,
       result: result,
       error: error,
       options: opts
     }}
  end

  defp build_timings(timings, total) do
    Enum.map(timings, fn {stage, us} ->
      %{
        stage: stage,
        elapsed_us: us,
        percentage: if(total > 0, do: us / total * 100, else: 0.0)
      }
    end)
  end

  defp format_time(us) when us < 1000, do: "#{us}μs"
  defp format_time(us) when us < 1_000_000, do: "#{Float.round(us / 1000, 1)}ms"
  defp format_time(us), do: "#{Float.round(us / 1_000_000, 2)}s"
end