# Benchmark script for Kiri Elixir NIF tokenizer performance.
#
# Usage:
#   cd packages/kiri && mix run ../../bench/tokenize.exs -- --dict core --corpus short
#
# Flags:
#   --dict small|core|full  Which dictionary size
#   --corpus short|long     Which corpus scale
#
# Outputs JSON to stdout with timing results.

target_duration_ms = 1000
dict_dir = Path.join(System.user_home!(), ".kiri-ji/dict")
bench_dir = Path.expand("../../bench", File.cwd!())

{opts, _, _} =
  OptionParser.parse(System.argv(), strict: [dict: :string, corpus: :string])

dict = Keyword.get(opts, :dict, "core")
corpus = Keyword.get(opts, :corpus, "short")
dict_path = Path.join(dict_dir, "system_#{dict}.dic")
kiri_dir = Path.join(System.user_home!(), ".kiri-ji/system_#{dict}.kiri")

unless File.exists?(kiri_dir) do
  IO.puts(:stderr, "Converted dictionary not found: #{kiri_dir}")
  IO.puts(:stderr, "Run: cd packages/kiri && mix kiri.convert #{dict_path} #{kiri_dir}")
  System.halt(1)
end

# Load corpus from JSON
corpus_path = Path.join([bench_dir, "corpus", "#{corpus}.json"])

unless File.exists?(corpus_path) do
  IO.puts(:stderr, "Corpus not found: #{corpus_path} — run: bun bench/gen-corpus.ts")
  System.halt(1)
end

texts =
  corpus_path
  |> File.read!()
  |> Jason.decode!()
  |> Enum.with_index(1)
  |> Enum.map(fn {text, i} ->
    name = "#{corpus}_#{String.pad_leading(Integer.to_string(i), 4, "0")}"
    {name, text}
  end)

corpus_chars_per_pass =
  Enum.reduce(texts, 0, fn {_name, text}, acc -> acc + String.length(text) end)

# Measure dictionary load time
dict_load_times =
  for _ <- 1..3 do
    {elapsed_us, {:ok, _d}} = :timer.tc(fn -> Kiri.load_dictionary(kiri_dir) end)
    elapsed_us / 1000
  end

dict_load_ms = Enum.sum(dict_load_times) / length(dict_load_times)

# Load dictionary for benchmarking
{:ok, tok} = Kiri.load_dictionary(kiri_dir)

# Calibration pass: run 1 iteration to estimate speed
{cal_us, _} =
  :timer.tc(fn ->
    for {_name, text} <- texts, do: Kiri.tokenize(tok, text)
  end)

cal_ms = cal_us / 1000
timed_iters = max(2, round(target_duration_ms / cal_ms))
warmup_iters = max(1, ceil(timed_iters / 10))

IO.puts(:stderr, "Corpus: #{length(texts)} texts, #{corpus_chars_per_pass} chars/pass, #{timed_iters} iterations (#{warmup_iters} warmup)")

# Warmup
for _ <- 1..warmup_iters do
  for {_name, text} <- texts, do: Kiri.tokenize(tok, text)
end

# Timed iterations
{elapsed_us, _} =
  :timer.tc(fn ->
    for _ <- 1..timed_iters do
      for {_name, text} <- texts, do: Kiri.tokenize(tok, text)
    end
  end)

total_ms = elapsed_us / 1000
total_chars = corpus_chars_per_pass * timed_iters
chars_per_sec = round(total_chars / (total_ms / 1000))

result = %{
  runtime: "elixir #{System.version()}",
  backend: "elixir",
  dict: dict,
  corpus: corpus,
  dictLoadMs: Float.round(dict_load_ms / 1.0, 1),
  totalChars: total_chars,
  timedIters: timed_iters,
  totalMs: Float.round(total_ms / 1.0, 1),
  charsPerSec: chars_per_sec
}

IO.puts(Jason.encode!(result, pretty: true))