# Benchmark: 3 versions of IEx.Autocomplete internals
# Run with: bin/elixir bench_autocomplete.exs

defmodule Bench do
  def measure(label, fun, iterations \\ 500) do
    # Warmup
    for _ <- 1..10, do: fun.()

    # Measure time
    :erlang.garbage_collect()
    {time_us, _} = :timer.tc(fn ->
      for _ <- 1..iterations, do: fun.()
    end)
    avg_us = time_us / iterations

    # Measure memory: total bytes allocated across N calls
    # Use binary_memory + heap tracking
    :erlang.garbage_collect()
    {:total_heap_size, heap0} = :erlang.process_info(self(), :total_heap_size)
    {:garbage_collection, gc_info0} = :erlang.process_info(self(), :garbage_collection)
    gcs0 = Keyword.get(gc_info0, :minor_gcs, 0)
    {:reductions, reds0} = :erlang.process_info(self(), :reductions)

    n = 50
    for _ <- 1..n, do: fun.()

    {:total_heap_size, heap1} = :erlang.process_info(self(), :total_heap_size)
    {:garbage_collection, gc_info1} = :erlang.process_info(self(), :garbage_collection)
    gcs1 = Keyword.get(gc_info1, :minor_gcs, 0)
    {:reductions, reds1} = :erlang.process_info(self(), :reductions)

    heap_words = div(heap1 - heap0, n)
    gcs = gcs1 - gcs0
    reds = div(reds1 - reds0, n)

    # Also measure with GC forced to see retained memory
    :erlang.garbage_collect()
    {:heap_size, retained0} = :erlang.process_info(self(), :heap_size)
    fun.()
    {:heap_size, retained1} = :erlang.process_info(self(), :heap_size)

    {avg_us, heap_words * 8, reds, gcs, (retained1 - retained0) * 8}
  end

  def compare(label, implementations, iterations \\ 500) do
    IO.puts(label)
    IO.puts(String.duplicate("-", 100))
    :io.format("  ~-30s ~10s ~12s ~10s ~8s ~12s~n",
      ["version", "time(µs)", "heap(bytes)", "reds", "GCs", "retained(B)"])
    :io.format("  ~-30s ~10s ~12s ~10s ~8s ~12s~n",
      ["-------", "--------", "-----------", "----", "---", "-----------"])

    results = for {name, fun} <- implementations do
      {avg_us, heap_bytes, reds, gcs, retained} = measure(name, fun, iterations)
      :io.format("  ~-30s ~10.1f ~12B ~10B ~8B ~12B~n",
        [name, avg_us, heap_bytes, reds, gcs, retained])
      {name, avg_us, heap_bytes}
    end

    # Show speedups relative to first
    [{base_name, base_time, base_mem} | rest] = results
    IO.puts("")
    for {name, time, mem} <- rest do
      time_x = if time > 0, do: Float.round(base_time / time, 1), else: :inf
      mem_x = if mem != 0 and base_mem != 0, do: Float.round(base_mem / mem, 1), else: "-"
      IO.puts("  #{name} vs #{base_name}: #{time_x}x faster, #{mem_x}x less heap")
    end
    IO.puts("")
  end
end

for mod <- [Enum, String, Map, Kernel, IO, File, Path, Agent, GenServer, Supervisor, Task,
            Logger, Macro, Module, Code, Regex, Range, Stream, URI, Access, Calendar,
            Date, DateTime, Time, NaiveDateTime, Keyword, List, Tuple, MapSet,
            Integer, Float, System, Process, Node, Port, Registry, DynamicSupervisor] do
  Code.ensure_loaded(mod)
end

loaded_count = length(:erlang.loaded())
IO.puts("=== IEx Autocomplete Benchmark — 3 versions (#{loaded_count} loaded modules) ===\n")

# ============================================================
# Version 1: BEFORE PR (:code.all_loaded + Enum.sort + dedup)
# ============================================================
before_get_modules = fn elixir_root? ->
  modules =
    if elixir_root? do
      ["Elixir.Elixir" | Enum.map(:code.all_loaded(), &Atom.to_string(elem(&1, 0)))]
    else
      Enum.map(:code.all_loaded(), &Atom.to_string(elem(&1, 0)))
    end

  case :code.get_mode() do
    :interactive ->
      apps = for [app] <- :ets.match(:ac_tab, {{:loaded, :"$1"}, :_}),
                 {:ok, mods} = :application.get_key(app, :modules),
                 mod <- mods, do: Atom.to_string(mod)
      modules ++ apps
    _ -> modules
  end
end

before_match_modules = fn hint, elixir_root? ->
  before_get_modules.(elixir_root?)
  |> Enum.sort()
  |> Enum.dedup()
  |> Enum.drop_while(&(not String.starts_with?(&1, hint)))
  |> Enum.take_while(&String.starts_with?(&1, hint))
end

before_match_elixir_modules = fn module, hint ->
  name = Atom.to_string(module)
  depth = length(String.split(name, ".")) + 1
  base = name <> "." <> hint

  for mod <- before_match_modules.(base, module == Elixir),
      parts = String.split(mod, "."),
      depth <= length(parts),
      name = Enum.at(parts, depth - 1),
      uniq: true,
      do: name
end

# ============================================================
# Version 2: CURRENT (:erlang.loaded + :lists.usort + binary_part)
# ============================================================
current_get_modules = fn elixir_root? ->
  modules =
    if elixir_root? do
      ["Elixir.Elixir" | Enum.map(:erlang.loaded(), &Atom.to_string/1)]
    else
      Enum.map(:erlang.loaded(), &Atom.to_string/1)
    end

  case :code.get_mode() do
    :interactive ->
      apps = for [app] <- :ets.match(:ac_tab, {{:loaded, :"$1"}, :_}),
                 {:ok, mods} = :application.get_key(app, :modules),
                 mod <- mods, do: Atom.to_string(mod)
      modules ++ apps
    _ -> modules
  end
end

current_match_modules = fn hint, elixir_root? ->
  current_get_modules.(elixir_root?)
  |> :lists.usort()
  |> Enum.drop_while(&(not String.starts_with?(&1, hint)))
  |> Enum.take_while(&String.starts_with?(&1, hint))
end

elixir_submodule_name = fn rest ->
  case :binary.match(rest, ".") do
    {pos, _} -> binary_part(rest, 0, pos)
    :nomatch -> rest
  end
end

current_match_elixir_modules = fn module, hint ->
  prefix = Atom.to_string(module) <> "."
  prefix_size = byte_size(prefix)
  base = prefix <> hint

  for mod <- current_match_modules.(base, module == Elixir),
      rest = binary_part(mod, prefix_size, byte_size(mod) - prefix_size),
      name = elixir_submodule_name.(rest),
      uniq: true,
      do: name
end

# ============================================================
# Version 3: PROPOSED (filter during collection, sort matches only)
# ============================================================
proposed_match_modules = fn hint, elixir_root? ->
  modules =
    if elixir_root? do
      acc = for mod <- :erlang.loaded(),
                str = Atom.to_string(mod),
                String.starts_with?(str, hint),
                do: str
      if String.starts_with?("Elixir.Elixir", hint), do: ["Elixir.Elixir" | acc], else: acc
    else
      for mod <- :erlang.loaded(),
          str = Atom.to_string(mod),
          String.starts_with?(str, hint),
          do: str
    end

  modules =
    case :code.get_mode() do
      :interactive ->
        apps = for [app] <- :ets.match(:ac_tab, {{:loaded, :"$1"}, :_}),
                   {:ok, mods} = :application.get_key(app, :modules),
                   mod <- mods,
                   str = Atom.to_string(mod),
                   String.starts_with?(str, hint),
                   do: str
        modules ++ apps
      _ -> modules
    end

  :lists.usort(modules)
end

proposed_match_elixir_modules = fn module, hint ->
  prefix = Atom.to_string(module) <> "."
  prefix_size = byte_size(prefix)
  base = prefix <> hint

  for mod <- proposed_match_modules.(base, module == Elixir),
      rest = binary_part(mod, prefix_size, byte_size(mod) - prefix_size),
      name = elixir_submodule_name.(rest),
      uniq: true,
      do: name
end

# ============================================================
# Verify correctness
# ============================================================
IO.puts("Verifying correctness...")
for {hint, root?} <- [{"Elixir.S", true}, {"Elixir.Enum", true}, {"er", false}, {"Elixir.", true}] do
  r1 = before_match_modules.(hint, root?)
  r2 = current_match_modules.(hint, root?)
  r3 = proposed_match_modules.(hint, root?)
  if r1 == r2 and r2 == r3 do
    IO.puts("  ✓ match_modules(#{inspect(hint)}) — #{length(r1)} results")
  else
    IO.puts("  ✗ match_modules(#{inspect(hint)}) MISMATCH: #{length(r1)} / #{length(r2)} / #{length(r3)}")
  end
end

for {mod, hint} <- [{Elixir, "S"}, {Elixir, "Enum"}, {Elixir, ""}] do
  r1 = before_match_elixir_modules.(mod, hint)
  r2 = current_match_elixir_modules.(mod, hint)
  r3 = proposed_match_elixir_modules.(mod, hint)
  if Enum.sort(r1) == Enum.sort(r2) and Enum.sort(r2) == Enum.sort(r3) do
    IO.puts("  ✓ match_elixir_modules(#{inspect(mod)}, #{inspect(hint)}) — #{length(r1)} results")
  else
    IO.puts("  ✗ match_elixir_modules(#{inspect(mod)}, #{inspect(hint)}) MISMATCH: #{length(r1)} / #{length(r2)} / #{length(r3)}")
    IO.puts("    before:   #{inspect(Enum.sort(r1), limit: 5)}")
    IO.puts("    current:  #{inspect(Enum.sort(r2), limit: 5)}")
    IO.puts("    proposed: #{inspect(Enum.sort(r3), limit: 5)}")
  end
end
IO.puts("")

# ============================================================
# Benchmarks
# ============================================================

Bench.compare("match_modules(\"Elixir.S\", true) — broad Elixir prefix", [
  {"v1: before PR", fn -> before_match_modules.("Elixir.S", true) end},
  {"v2: current", fn -> current_match_modules.("Elixir.S", true) end},
  {"v3: proposed", fn -> proposed_match_modules.("Elixir.S", true) end},
])

Bench.compare("match_modules(\"Elixir.Enum\", true) — specific module", [
  {"v1: before PR", fn -> before_match_modules.("Elixir.Enum", true) end},
  {"v2: current", fn -> current_match_modules.("Elixir.Enum", true) end},
  {"v3: proposed", fn -> proposed_match_modules.("Elixir.Enum", true) end},
])

Bench.compare("match_modules(\"Elixir.\", true) — all Elixir modules (worst case)", [
  {"v1: before PR", fn -> before_match_modules.("Elixir.", true) end},
  {"v2: current", fn -> current_match_modules.("Elixir.", true) end},
  {"v3: proposed", fn -> proposed_match_modules.("Elixir.", true) end},
])

Bench.compare("match_modules(\"er\", false) — erlang prefix", [
  {"v1: before PR", fn -> before_match_modules.("er", false) end},
  {"v2: current", fn -> current_match_modules.("er", false) end},
  {"v3: proposed", fn -> proposed_match_modules.("er", false) end},
])

Bench.compare("match_elixir_modules(Elixir, \"S\") — full submodule extraction", [
  {"v1: before PR", fn -> before_match_elixir_modules.(Elixir, "S") end},
  {"v2: current", fn -> current_match_elixir_modules.(Elixir, "S") end},
  {"v3: proposed", fn -> proposed_match_elixir_modules.(Elixir, "S") end},
])

Bench.compare("match_elixir_modules(Elixir, \"\") — all modules (tab after .)", [
  {"v1: before PR", fn -> before_match_elixir_modules.(Elixir, "") end},
  {"v2: current", fn -> current_match_elixir_modules.(Elixir, "") end},
  {"v3: proposed", fn -> proposed_match_elixir_modules.(Elixir, "") end},
])

IO.puts("=== Done ===")