# frozen_string_literal: true
require "net/http"
require "json"
URI_ = URI("http://127.0.0.1:11434/api/generate")
def one(prompt)
  t = Process.clock_gettime(Process::CLOCK_MONOTONIC)
  res = Net::HTTP.start(URI_.host, URI_.port, read_timeout: 600) do |h|
    h.post(URI_.path, { model: "qwen3:4b", prompt: prompt, stream: false, think: false,
                       options: { num_predict: 200, temperature: 0, seed: 1 } }.to_json, "Content-Type" => "application/json")
  end
  j = JSON.parse(res.body)
  { wall: Process.clock_gettime(Process::CLOCK_MONOTONIC) - t,
    prompt_tokens: j["prompt_eval_count"], prefill_tps: j["prompt_eval_count"].to_f / (j["prompt_eval_duration"] / 1e9),
    gen_tokens: j["eval_count"], decode_tps: j["eval_count"].to_f / (j["eval_duration"] / 1e9) }
end
prompt = "Explain in plain words how a key-value cache speeds up text generation in a transformer. " * 4
one("hi") # warm
ARGV.map(&:to_i).each do |n|
  t = Process.clock_gettime(Process::CLOCK_MONOTONIC)
  rs = Array.new(n) { Thread.new { one(prompt) } }.map(&:value)
  wall = Process.clock_gettime(Process::CLOCK_MONOTONIC) - t
  puts({ concurrency: n, wall_s: wall.round(2), total_gen_tokens: rs.sum { _1[:gen_tokens] },
         aggregate_tps: (rs.sum { _1[:gen_tokens] } / wall).round(1),
         per_req_decode_tps: rs.map { _1[:decode_tps].round(1) }, prefill_tps: rs.map { _1[:prefill_tps].round(0) },
         max_latency_s: rs.map { _1[:wall] }.max.round(2) }.to_json)
end