# Concurrent Elixir benchmark: distributes corpus across N workers.
#
# Usage:
#   cd packages/kiri && mix run ../../bench/tokenize_concurrent.exs -- --dict core --corpus short
#
# Flags:
#   --dict small|core|full    Which dictionary size
#   --corpus short|long       Which corpus scale
#   --concurrency N           Number of concurrent workers (default: 2048)
#   --backend elixir|nif      Which backend to use (default: elixir)
#
# Outputs JSON to stdout with timing results.

bench_dir = Path.expand("../../bench", File.cwd!())

# mix run passes "--" as first element of System.argv(); strip it so
# OptionParser doesn't treat it as end-of-options.
argv =
  case System.argv() do
    ["--" | rest] -> rest
    other -> other
  end

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

dict_name = Keyword.get(opts, :dict, "core")
corpus = Keyword.get(opts, :corpus, "short")
concurrency = Keyword.get(opts, :concurrency, 2048)
backend_name = Keyword.get(opts, :backend, "elixir")

backend_opt =
  case backend_name do
    "nif" -> [backend: :nif]
    "elixir" -> [backend: :elixir]
    _ ->
      IO.puts(:stderr, "Error: --backend must be elixir or nif")
      System.halt(1)
  end
kiri_dir = Path.join(System.user_home!(), ".kiri-ji/system_#{dict_name}.kiri")
dic_path = Path.join(System.user_home!(), ".kiri-ji/dict/system_#{dict_name}.dic")
load_opts = if File.exists?(dic_path), do: [dic_path: dic_path], else: []

unless File.exists?(kiri_dir) do
  IO.puts(:stderr, "Converted dictionary not found: #{kiri_dir}")
  System.halt(1)
end

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

unless File.exists?(corpus_path) do
  IO.puts(:stderr, "Corpus not found: #{corpus_path}")
  System.halt(1)
end

texts = corpus_path |> File.read!() |> Jason.decode!()
corpus_chars = Enum.reduce(texts, 0, fn text, acc -> acc + String.length(text) end)

IO.puts(
  :stderr,
  "Dict: #{dict_name}, Corpus: #{corpus} (#{length(texts)} texts, #{corpus_chars} chars)"
)

IO.puts(:stderr, "Backend: #{backend_name}, Concurrency: #{concurrency}, Schedulers: #{System.schedulers_online()}")

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

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

{:ok, dict} = Kiri.load_dictionary(kiri_dir, load_opts)
IO.puts(:stderr, "Dict loaded (avg #{Float.round(dict_load_ms, 1)}ms)")

# Assign texts per worker: 10 for short corpus, 1 for long
texts_per_worker = if corpus == "short", do: 10, else: 1
num_texts = length(texts)

assignments =
  for i <- 0..(concurrency - 1) do
    for j <- 0..(texts_per_worker - 1) do
      Enum.at(texts, rem(i * texts_per_worker + j, num_texts))
    end
  end

total_chars_processed =
  assignments |> List.flatten() |> Enum.reduce(0, fn text, acc -> acc + String.length(text) end)

IO.puts(
  :stderr,
  "#{concurrency} workers × #{texts_per_worker} texts each, #{num_texts} unique texts"
)

# Warmup: run a subset of workers to warm up schedulers
assignments
|> Enum.take(min(concurrency, num_texts))
|> Enum.map(fn batch ->
  Task.async(fn -> Enum.each(batch, fn text -> Kiri.tokenize(dict, text, backend_opt) end) end)
end)
|> Task.await_many(:infinity)

IO.puts(:stderr, "Warmup done, running timed pass...")

# Timed: each worker tokenizes its batch of texts
{elapsed_us, _} =
  :timer.tc(fn ->
    assignments
    |> Enum.map(fn batch ->
      Task.async(fn -> Enum.each(batch, fn text -> Kiri.tokenize(dict, text, backend_opt) end) end)
    end)
    |> Task.await_many(:infinity)
  end)

total_ms = elapsed_us / 1000
chars_per_sec = round(total_chars_processed / (total_ms / 1000))

IO.puts(:stderr, "Done: #{Float.round(total_ms / 1000, 1)}s, #{chars_per_sec} chars/s")

result = %{
  runtime: "elixir #{System.version()}",
  backend: backend_name,
  dict: dict_name,
  corpus: corpus,
  concurrency: concurrency,
  textsPerWorker: texts_per_worker,
  schedulers: System.schedulers_online(),
  dictLoadMs: Float.round(dict_load_ms, 1),
  totalChars: total_chars_processed,
  totalMs: Float.round(total_ms, 1),
  charsPerSec: chars_per_sec
}

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