defmodule Mix.Tasks.Firebird.Bench do
  @moduledoc """
  Benchmarks comparing compiled WASM vs native BEAM execution.

  ## Usage

      mix firebird.bench
      mix firebird.bench --functions fibonacci,factorial
      mix firebird.bench --iterations 1000
      mix firebird.bench --warmup 20
      mix firebird.bench --output bench_results.json
      mix firebird.bench --format markdown
  """

  use Mix.Task

  @shortdoc "Benchmark WASM vs native BEAM execution"

  @switches [
    functions: :string,
    iterations: :integer,
    warmup: :integer,
    output: :string,
    format: :string,
    verbose: :boolean,
    compare: :string
  ]

  @aliases [
    f: :functions,
    n: :iterations,
    o: :output,
    v: :verbose,
    w: :warmup
  ]

  @impl Mix.Task
  def run(args) do
    {opts, _, _} = OptionParser.parse(args, switches: @switches, aliases: @aliases)

    verbose = Keyword.get(opts, :verbose, false)
    format = Keyword.get(opts, :format, "table")

    Mix.shell().info("🔥 Firebird WASM vs BEAM Benchmark")
    Mix.shell().info("")

    benchmarks = build_benchmarks(opts)

    results =
      Enum.map(benchmarks, fn bench ->
        run_benchmark(bench, verbose)
      end)

    # Display results
    case format do
      "json" -> Mix.shell().info(Firebird.Benchmark.format_json(results))
      "csv" -> Mix.shell().info(Firebird.Benchmark.format_csv(results))
      "markdown" -> Mix.shell().info(Firebird.Benchmark.format_markdown(results))
      _ -> print_results(results)
    end

    # Compare with previous results if requested
    case Keyword.get(opts, :compare) do
      nil -> :ok
      path -> compare_with_previous(results, path)
    end

    if output = Keyword.get(opts, :output) do
      save_results(results, output)
    end

    :ok
  end

  @doc """
  Run a benchmark comparing WASM and BEAM execution.

  Returns a map with timing results for both execution modes.
  """
  @spec run_benchmark(map(), boolean()) :: map()
  def run_benchmark(bench, verbose \\ false) do
    %{name: name, source: source, function: function, args_list: args_list} = bench

    if verbose do
      Mix.shell().info("  Benchmarking #{name}...")
    end

    iterations = bench[:iterations] || 100
    warmup = bench[:warmup] || 10

    result =
      Firebird.Benchmark.compare(
        source,
        function,
        args_list,
        beam_fun: bench.beam_fun,
        iterations: iterations,
        warmup: warmup
      )

    Map.put(result, :name, name)
  end

  defp build_benchmarks(opts) do
    requested =
      case Keyword.get(opts, :functions) do
        nil -> nil
        str -> String.split(str, ",", trim: true) |> Enum.map(&String.trim/1)
      end

    iterations = Keyword.get(opts, :iterations, 100)
    warmup = Keyword.get(opts, :warmup, 10)

    all_benchmarks = [
      %{
        name: "add (simple arithmetic)",
        source: """
        defmodule BenchAdd do
          @wasm true
          def add(a, b), do: a + b
        end
        """,
        function: :add,
        beam_fun: fn a, b -> a + b end,
        args_list: for(a <- 1..20, b <- 1..20, do: [a, b]),
        iterations: iterations,
        warmup: warmup
      },
      %{
        name: "fibonacci(20) (recursive)",
        source: """
        defmodule BenchFib do
          @wasm true
          def fibonacci(0), do: 0
          def fibonacci(1), do: 1
          def fibonacci(n), do: fibonacci(n - 1) + fibonacci(n - 2)
        end
        """,
        function: :fibonacci,
        beam_fun: &beam_fibonacci/1,
        args_list: [[20]],
        iterations: iterations,
        warmup: warmup
      },
      %{
        name: "factorial(15) (recursive)",
        source: """
        defmodule BenchFact do
          @wasm true
          def factorial(0), do: 1
          def factorial(n), do: n * factorial(n - 1)
        end
        """,
        function: :factorial,
        beam_fun: &beam_factorial/1,
        args_list: [[15]],
        iterations: iterations,
        warmup: warmup
      },
      %{
        name: "gcd (Euclidean)",
        source: """
        defmodule BenchGcd do
          @wasm true
          def gcd(a, 0), do: a
          def gcd(a, b), do: gcd(b, rem(a, b))
        end
        """,
        function: :gcd,
        beam_fun: &beam_gcd/2,
        args_list: for(a <- [12, 100, 252, 1071], b <- [8, 75, 105, 462], do: [a, b]),
        iterations: iterations,
        warmup: warmup
      },
      %{
        name: "sum_to(100) (recursive sum)",
        source: """
        defmodule BenchSum do
          @wasm true
          def sum_to(0), do: 0
          def sum_to(n), do: n + sum_to(n - 1)
        end
        """,
        function: :sum_to,
        beam_fun: &beam_sum_to/1,
        args_list: [[100]],
        iterations: iterations,
        warmup: warmup
      },
      %{
        name: "power(2, 20) (exponentiation)",
        source: """
        defmodule BenchPow do
          @wasm true
          def power(_, 0), do: 1
          def power(base, exp), do: base * power(base, exp - 1)
        end
        """,
        function: :power,
        beam_fun: &beam_power/2,
        args_list: [[2, 20], [3, 10], [5, 8]],
        iterations: iterations,
        warmup: warmup
      },
      %{
        name: "collatz_steps(27)",
        source: """
        defmodule BenchCollatz do
          @wasm true
          def collatz_steps(1), do: 0
          def collatz_steps(n) do
            if rem(n, 2) == 0 do
              1 + collatz_steps(div(n, 2))
            else
              1 + collatz_steps(3 * n + 1)
            end
          end
        end
        """,
        function: :collatz_steps,
        beam_fun: &beam_collatz_steps/1,
        args_list: [[27]],
        iterations: iterations,
        warmup: warmup
      },
      %{
        name: "is_even check",
        source: """
        defmodule BenchIsEven do
          @wasm true
          def is_even(n) do
            if rem(n, 2) == 0 do
              1
            else
              0
            end
          end
        end
        """,
        function: :is_even,
        beam_fun: fn n -> if rem(n, 2) == 0, do: 1, else: 0 end,
        args_list: for(n <- 1..100, do: [n]),
        iterations: iterations,
        warmup: warmup
      }
    ]

    if requested do
      Enum.filter(all_benchmarks, fn b ->
        Atom.to_string(b.function) in requested or b.name in requested
      end)
    else
      all_benchmarks
    end
  end

  # BEAM reference implementations
  defp beam_fibonacci(0), do: 0
  defp beam_fibonacci(1), do: 1
  defp beam_fibonacci(n), do: beam_fibonacci(n - 1) + beam_fibonacci(n - 2)

  defp beam_factorial(0), do: 1
  defp beam_factorial(n), do: n * beam_factorial(n - 1)

  defp beam_gcd(a, 0), do: a
  defp beam_gcd(a, b), do: beam_gcd(b, rem(a, b))

  defp beam_sum_to(0), do: 0
  defp beam_sum_to(n), do: n + beam_sum_to(n - 1)

  defp beam_power(_, 0), do: 1
  defp beam_power(base, exp), do: base * beam_power(base, exp - 1)

  defp beam_collatz_steps(1), do: 0
  defp beam_collatz_steps(n) when rem(n, 2) == 0, do: 1 + beam_collatz_steps(div(n, 2))
  defp beam_collatz_steps(n), do: 1 + beam_collatz_steps(3 * n + 1)

  defp print_results(results) do
    Mix.shell().info(
      "┌─────────────────────────────────────┬───────────┬───────────┬──────────┬──────────┬──────────┐"
    )

    Mix.shell().info(
      "│ Benchmark                           │ WASM (μs) │ BEAM (μs) │ Speedup  │ WASM p95 │ BEAM p95 │"
    )

    Mix.shell().info(
      "├─────────────────────────────────────┼───────────┼───────────┼──────────┼──────────┼──────────┤"
    )

    for r <- results do
      name = String.pad_trailing(r.name, 37) |> String.slice(0, 37)
      wasm = r.wasm.avg_us |> Float.round(1) |> to_string() |> String.pad_leading(9)
      beam = r.beam.avg_us |> Float.round(1) |> to_string() |> String.pad_leading(9)
      wasm_p95 = r.wasm.p95_us |> Float.round(1) |> to_string() |> String.pad_leading(8)
      beam_p95 = r.beam.p95_us |> Float.round(1) |> to_string() |> String.pad_leading(8)

      speedup_str =
        if r.speedup >= 1.0 do
          "#{Float.round(r.speedup, 2)}x"
        else
          inv = Float.round(1.0 / max(r.speedup, 0.001), 2)
          "#{inv}x slow"
        end

      speedup_str = String.pad_leading(speedup_str, 8)

      Mix.shell().info(
        "│ #{name} │ #{wasm} │ #{beam} │ #{speedup_str} │ #{wasm_p95} │ #{beam_p95} │"
      )
    end

    Mix.shell().info(
      "└─────────────────────────────────────┴───────────┴───────────┴──────────┴──────────┴──────────┘"
    )

    Mix.shell().info("")
    Mix.shell().info("Note: WASM times include Wasmex FFI overhead. Raw WASM execution")
    Mix.shell().info("is typically faster but measured end-to-end from Elixir.")
  end

  defp compare_with_previous(current_results, path) do
    case File.read(path) do
      {:ok, json} ->
        case Jason.decode(json) do
          {:ok, previous} ->
            Mix.shell().info("")
            Mix.shell().info("📊 Comparison with previous results (#{path}):")

            Enum.each(current_results, fn curr ->
              name = curr.name

              prev =
                Enum.find(previous, fn p ->
                  (is_map(p) and p["name"] == name) or
                    (is_map(p) and Map.get(p, :name) == name)
                end)

              if prev do
                prev_wasm =
                  get_in(prev, ["wasm", "avg_us"]) || get_in(prev, [:wasm, :avg_us]) || 0

                curr_wasm = curr.wasm.avg_us

                if prev_wasm > 0 do
                  change = (curr_wasm - prev_wasm) / prev_wasm * 100
                  indicator = if change < 0, do: "✅", else: "⚠"
                  Mix.shell().info("   #{indicator} #{name}: #{Float.round(change, 1)}%")
                end
              end
            end)

          _ ->
            Mix.shell().error("   ⚠ Could not parse previous results from #{path}")
        end

      {:error, _} ->
        Mix.shell().error("   ⚠ Previous results file not found: #{path}")
    end
  end

  defp save_results(results, path) do
    json = Jason.encode!(results, pretty: true)
    File.write!(path, json)
    Mix.shell().info("📊 Results saved to #{path}")
  end
end