-- small_load.lua — hard 2500 rps: hot=2000, cold=500 (time-driven dual schedulers)

local ffi = require("ffi")
ffi.cdef[[
  typedef long time_t;
  struct timeval { time_t tv_sec; time_t tv_usec; };
  int gettimeofday(struct timeval* tv, void* tz);
]]
local tv = ffi.new("struct timeval")
local function mono_us()
  ffi.C.gettimeofday(tv, nil)
  return tonumber(tv.tv_sec) * 1000000 + tonumber(tv.tv_usec)
end

local function getenv(k, def) local v=os.getenv(k); if v==nil or v=="" then return def end; return v end
local function getenv_num(k, def) local v=os.getenv(k); if v==nil or v=="" then return def end; v=tonumber(v); if not v then return def end; return v end

-- targets (can be overridden via ENV)
local TARGET_RPS = getenv_num("TARGET_RPS", 2500)
local HOT_RPS    = getenv_num("HOT_RPS",    2000)
local COLD_RPS   = getenv_num("COLD_RPS",    500)
local HOT_RING   = math.max(1, getenv_num("HOT_RING", 1000000))
local CPT        = math.max(1, getenv_num("CPT", 1))
local ENDPOINT   = getenv("ENDPOINT", "/api/v2/pagedata")
local HOSTHDR    = getenv("HOST", "")

-- normalize if needed
do
  local sum = HOT_RPS + COLD_RPS
  if sum ~= TARGET_RPS then
    HOT_RPS  = math.floor(TARGET_RPS * (HOT_RPS  / sum) + 0.5)
    COLD_RPS = TARGET_RPS - HOT_RPS
  end
end

-- per-thread state
local threads = {}
local THREADS = 0

local function encode_brackets(s)
  s = string.gsub(s, "%[", "%%5B"); s = string.gsub(s, "%]", "%%5D"); return s
end

local function build_query(i)
  local q =
    "?project[id]=285&domain=example.com&language=en" ..
    "&choice[name]=betting" ..
    "&choice[choice][name]=betting_live" ..
    "&choice[choice][choice][name]=betting_live_null" ..
    "&choice[choice][choice][choice][name]=betting_live_null_"..i ..
    "&choice[choice][choice][choice][choice][name]=betting_live_null_"..i.."_"..i ..
    "&choice[choice][choice][choice][choice][choice][name]=betting_live_null_"..i.."_"..i.."_"..i ..
    "&choice[choice][choice][choice][choice][choice][choice]=null"
  return encode_brackets(q)
end

local function seed_from_thread(thread)
  local s=tostring(thread); local h=0
  for i=1,#s do h=(h*131 + string.byte(s,i)) % 2147483647 end
  local tsec=os.time()%2147483647; local noise=math.floor((os.clock()%1)*1e6)
  return (h + tsec + noise) % 2147483647
end

local function hdr_hot()
  if HOSTHDR ~= "" then
    return {
      ["Host"] = HOSTHDR,
      ["Accept-Encoding"] = "gzip",
      ["Accept-Language"] = "en-US,en;q=0.9",
      ["User-Agent"]      = "wrk-2500rps/hot",
      ["Content-Type"]    = "application/json",
    }
  end
  return {
    ["Accept-Encoding"] = "gzip",
    ["Accept-Language"] = "en-US,en;q=0.9",
    ["User-Agent"]      = "wrk-2500rps/hot",
    ["Content-Type"]    = "application/json",
  }
end

local function hdr_cold()
  if HOSTHDR ~= "" then
    return {
      ["Host"] = HOSTHDR,
      ["Accept-Encoding"] = "br",
      ["Accept-Language"] = "en-US,en;q=0.9",
      ["User-Agent"]      = "wrk-2500rps/cold",
      ["Content-Type"]    = "application/json",
    }
  end
  return {
    ["Accept-Encoding"] = "br",
    ["Accept-Language"] = "en-US,en;q=0.9",
    ["User-Agent"]      = "wrk-2500rps/cold",
    ["Content-Type"]    = "application/json",
  }
end

-- next keys
local function next_hot(st)
  st.hot_idx = (st.hot_idx + 1) % st.hot_ring
  return st.hot_base + st.hot_idx
end
local function next_cold(st)
  st.cold_seq = st.cold_seq + 1
  return st.cold_base + st.cold_seq
end

-- per-thread rates → periods
local function per_thread_periods()
  local thr = math.max(1, THREADS)
  local hot_rps  = HOT_RPS  / thr / CPT
  local cold_rps = COLD_RPS / thr / CPT
  if hot_rps  < 0.0001 then hot_rps  = 0.0001 end
  if cold_rps < 0.0001 then cold_rps = 0.0001 end
  return math.floor(1e6 / hot_rps + 0.5), math.floor(1e6 / cold_rps + 0.5)
end

local function ensure_state()
  local t = wrk.thread
  local st = threads[t]
  if st then return st end

  math.randomseed(seed_from_thread(t))
  local Th, Tc = per_thread_periods()
  local now = mono_us()

  st = {
    -- rings
    hot_ring   = HOT_RING,
    hot_idx    = math.random(0, math.max(1,HOT_RING)-1),
    hot_base   = 1,
    -- cold stream (disjoint per thread)
    cold_base  = 10^9 + 1000000 * (THREADS + 1),
    cold_seq   = 0,
    -- headers
    h_hdr      = hdr_hot(),
    c_hdr      = hdr_cold(),
    -- schedulers
    Th         = Th,          -- hot period, us
    Tc         = Tc,          -- cold period, us
    next_hot   = now + Th,
    next_cold  = now + Tc,
    last_recalc= now,
  }
  threads[t] = st
  return st
end

function setup(thread)
  THREADS = THREADS + 1
  math.randomseed(seed_from_thread(thread))

  local Th, Tc = per_thread_periods()
  local now = mono_us()
  threads[thread] = {
    hot_ring   = HOT_RING,
    hot_idx    = math.random(0, math.max(1,HOT_RING)-1),
    hot_base   = 1,
    cold_base  = 10^9 + 1000000 * THREADS,
    cold_seq   = 0,
    h_hdr      = hdr_hot(),
    c_hdr      = hdr_cold(),
    Th         = Th,
    Tc         = Tc,
    next_hot   = now + Th,
    next_cold  = now + Tc,
    last_recalc= now,
  }
end

function init(args)
  -- nothing; guard with ensure_state() in delay/request
end

-- helper: advance timer in fixed quanta so we don't "burst"
local function advance_due(due, period, now)
  if now <= due then return due end
  -- catch-up: hop by whole periods until due > now
  local missed = math.floor((now - due) / period) + 1
  return due + missed * period
end

function delay()
  local st = ensure_state()
  local now = mono_us()

  -- re-calc per second (if thread count changed or CPT tuned)
  if now - st.last_recalc >= 1000000 then
    local Th, Tc = per_thread_periods()
    st.Th, st.Tc = Th, Tc
    -- re-anchor timers against 'now' to avoid drift
    st.next_hot  = advance_due(st.next_hot,  st.Th, now)
    st.next_cold = advance_due(st.next_cold, st.Tc, now)
    st.last_recalc = now
  end

  local next_due = st.next_hot
  if st.next_cold < next_due then next_due = st.next_cold end
  local d = next_due - now
  if d < 0 then d = 0 end
  return d
end

request = function()
  local st = ensure_state()
  local now = mono_us()

  -- choose which stream is due first
  local is_hot = st.next_hot <= st.next_cold
  if is_hot then
    -- schedule next hot tick
    st.next_hot = advance_due(st.next_hot, st.Th, now)
    local k = next_hot(st)
    return wrk.format("GET", ENDPOINT .. build_query(k), st.h_hdr)
  else
    st.next_cold = advance_due(st.next_cold, st.Tc, now)
    local k = next_cold(st)
    return wrk.format("GET", ENDPOINT .. build_query(k), st.c_hdr)
  end
end

function done(summary, latency, requests)
  io.write(string.format(
    "\n[RPS] target=%d hot=%d cold=%d  ring=%d  threads=%d CPT=%d\n",
    TARGET_RPS, HOT_RPS, COLD_RPS, HOT_RING, THREADS, CPT))
  io.write(string.format(
    "[Latency] mean=%.2fms p99=%.2fms max=%.2fms  req/s=%.0f\n\n",
    latency.mean/1000, latency:percentile(99)/1000, latency.max/1000,
    summary.requests / (summary.duration/1000000)))
end
