-- stress_admission.lua — LFU-proof, infinite keyspace, non-blocking

-- ========== ENV ==========
-- MODE       = lfu | avalanche | mixed | uniform (default: lfu)
-- PHASE_MODE = ops | time (default: ops)
--   ops : OPS_WARM / OPS_BATTLE / OPS_SHIFT
--   time: WARM_SEC / BATTLE_SEC / SHIFT_SEC  (wall clock via os.time())
-- ENDPOINTS  = comma-separated paths (default below)
-- PATH_MODE  = fixed | rr | random (default: fixed)
-- PATH_FIXED = index for fixed (1-based, default 1)
-- HOT_RING   = number of A/B pairs (default 4000)
-- OPS_WARM=200000 OPS_BATTLE=300000 OPS_SHIFT=300000
-- WARM_SEC=20 BATTLE_SEC=30 SHIFT_SEC=30
-- CHALLENGE_EVERY=10          -- battle: how often try B (with double-touch)
-- SHIFT_CHALLENGE_EVERY=10    -- shift : how often try A (with double-touch)
-- AE = identity|gzip|br (default: identity)
-- HOST = host:port (optional)
-- LOG_PHASE = 0|1 (default 1)
--
-- Avalanche mode:
--   MODE=avalanche  → бесконечный поток уникальных ключей (чистый холод).
-- =========================

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

local function split_csv(s)
  local t = {}
  for part in string.gmatch(s or "", "([^,]+)") do
    local x = string.gsub(part, "^%s*(.-)%s*$", "%1")
    if x ~= "" then table.insert(t, x) end
  end
  return t
end

-- -------- config --------
local MODE            = getenv("MODE","lfu")                -- lfu | avalanche | mixed | uniform
local PHASE_MODE      = getenv("PHASE_MODE","ops")          -- ops | time
local ENDPOINTS       = split_csv(getenv("ENDPOINTS",
  "/api/v2/pagedata,/api/v2/pagecontent,/api/v1/visual,/api/v1/micro-markup"))
if #ENDPOINTS == 0 then ENDPOINTS = {"/api/v2/pagedata"} end

local PATH_MODE       = getenv("PATH_MODE","fixed")
local PATH_FIXED      = getenv_num("PATH_FIXED",1)

local HOT_RING        = math.max(1, getenv_num("HOT_RING", 1250000))

local OPS_WARM        = math.max(1, getenv_num("OPS_WARM",   600000))
local OPS_BATTLE      = math.max(1, getenv_num("OPS_BATTLE", 325000))
local OPS_SHIFT       = math.max(1, getenv_num("OPS_SHIFT",  325000))

local WARM_SEC        = math.max(1, getenv_num("WARM_SEC",   20))
local BATTLE_SEC      = math.max(1, getenv_num("BATTLE_SEC", 30))
local SHIFT_SEC       = math.max(1, getenv_num("SHIFT_SEC",  30))

local CH_EVERY        = math.max(1, getenv_num("CHALLENGE_EVERY",         10))
local SHIFT_CH_EVERY  = math.max(1, getenv_num("SHIFT_CHALLENGE_EVERY",   10))

local AE_HEADER       = getenv("AE","identity")
local HOSTHDR         = getenv("HOST","")
local LOG_PHASE       = getenv_num("LOG_PHASE",1)

-- -------- shared/thread state --------
local START_TS = nil     -- wall clock epoch (os.time())
local threads  = {}

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

local function build_query(i)
  -- 'i' is integer (<= 9e15 for Lua number safety)
  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 pick_endpoint(st)
  local n = #ENDPOINTS
  if PATH_MODE == "fixed" then
    local idx = PATH_FIXED; if idx < 1 then idx = 1 end; if idx > n then idx = n end
    return ENDPOINTS[idx]
  elseif PATH_MODE == "rr" then
    st.ep_rr = (st.ep_rr + 1) % n
    return ENDPOINTS[st.ep_rr + 1]
  else
    local idx = 1 + math.random(0, n-1)
    return ENDPOINTS[idx]
  end
end

-- LFU rings A/B
local function key_A(j) return j + 1 end
local function key_B(j) return HOT_RING + j + 1 end
local function next_A(st) st.a_idx = (st.a_idx + 1) % st.hot_ring; return st.a_idx end
local function next_B(st) st.b_idx = (st.b_idx + 1) % st.hot_ring; return st.b_idx end

-- Infinite unique stream (per-thread disjoint):
-- 53-bit safe region for Lua numbers: < 9e15. Строим базу и шагаем.
local UNIQUE_STRIDE = 1000000  -- разнесение потоков
local function next_unique(st)
  st.u_seq = st.u_seq + 1
  return st.u_base + st.u_seq
end

local function build_request(path, i, phase)
  local full = path .. build_query(i)
  local hdrs = {
    ["Accept-Encoding"]=AE_HEADER,
    ["Accept-Language"]="en-US,en;q=0.9",
    ["Content-Type"]="application/json",
    ["User-Agent"]="wrk-tinylfu/lfu-infinite-1.0",
    ["X-LFU-Phase"]=phase or "n/a",
  }
  if HOSTHDR ~= "" then hdrs["Host"] = HOSTHDR end
  return wrk.format("GET", full, hdrs)
end

-- -------- фазирование --------
local function phase_by_ops(st)
  local n = st.req_total
  if n < OPS_WARM then return "warm"
  elseif n < OPS_WARM + OPS_BATTLE then return "battle"
  elseif n < OPS_WARM + OPS_BATTLE + OPS_SHIFT then return "shift"
  else return "shift" end
end

local function phase_by_time(now)
  local t1 = START_TS + WARM_SEC
  local t2 = t1 + BATTLE_SEC
  local t3 = t2 + SHIFT_SEC
  if now <= t1 then return "warm"
  elseif now <= t2 then return "battle"
  elseif now <= t3 then return "shift"
  else return "shift" end
end

local function now_phase(st)
  if MODE ~= "lfu" then return MODE end
  if PHASE_MODE == "time" then
    local now = os.time()  -- wall clock seconds
    return phase_by_time(now)
  else
    return phase_by_ops(st)
  end
end

local function log_phase_start(st, phase)
  if LOG_PHASE == 0 then return end
  if st.last_phase ~= phase then
    st.last_phase = phase
    io.write(string.format("[wrk-lfu] thread=%s phase=%s start req=%d t=%ds path=%s AE=%s\n",
      tostring(wrk.thread), phase, st.req_total,
      (START_TS and (os.time()-START_TS) or 0), PATH_MODE, AE_HEADER))
  end
end

-- -------- wrk lifecycle --------
function setup(thread)
  if START_TS == nil then START_TS = os.time() end
  math.randomseed(seed_from_thread(thread))
  local st = {
    -- rings
    hot_ring=HOT_RING, a_idx=math.random(0, math.max(1,HOT_RING)-1), b_idx=0,
    -- endpoints
    ep_rr=0,
    -- counters
    req_total=0, a_hits=0, b_hits=0, a_chal=0, b_chal=0, repeats=0,
    -- phase
    last_phase="",
    -- DK double-touch
    pending_repeat=false, repeat_key=0,
    -- cadences
    chal_every=CH_EVERY, shift_every=SHIFT_CH_EVERY,
    -- unique stream (disjoint per thread)
    u_base = 10^9 + UNIQUE_STRIDE * (math.random(1, 9000)), -- разнесение потоков
    u_seq  = 0,
  }
  threads[thread] = st
end

function init(args) end

-- безопасный генератор: всегда возвращает строку запроса
local function generate_request(st)
  st.req_total = st.req_total + 1
  -- avalanche mode: бесконечный поток новых ключей
  if MODE == "avalanche" then
    local k = next_unique(st)
    local p = pick_endpoint(st); return build_request(p, k, "avalanche")
  end

  local phase = now_phase(st)
  log_phase_start(st, phase)

  -- DK 2nd touch
  if st.pending_repeat then
    st.pending_repeat = false
    st.repeats = st.repeats + 1
    local p = pick_endpoint(st); return build_request(p, st.repeat_key, phase)
  end

  if MODE == "lfu" then
    if phase == "warm" then
      local j = next_A(st); st.a_hits = st.a_hits + 1
      local p = pick_endpoint(st); return build_request(p, key_A(j), phase)

    elseif phase == "battle" then
      local do_ch = (st.req_total % st.chal_every == 0)
      if do_ch then
        local j = next_B(st); local k = key_B(j)
        st.b_chal = st.b_chal + 1
        st.pending_repeat = true; st.repeat_key = k
        local p = pick_endpoint(st); return build_request(p, k, phase)
      else
        local j = next_A(st); st.a_hits = st.a_hits + 1
        local p = pick_endpoint(st); return build_request(p, key_A(j), phase)
      end

    else -- shift
      local do_ch = (st.req_total % st.shift_every == 0)
      if do_ch then
        local j = next_A(st); local k = key_A(j)
        st.a_chal = st.a_chal + 1
        st.pending_repeat = true; st.repeat_key = k
        local p = pick_endpoint(st); return build_request(p, k, phase)
      else
        local j = next_B(st); st.b_hits = st.b_hits + 1
        local p = pick_endpoint(st); return build_request(p, key_B(j), phase)
      end
    end

  elseif MODE == "mixed" then
    -- «Zipf-ish»: чаще A, иногда уникальные
    if st.req_total % 20 == 0 then
      local k = next_unique(st)
      local p = pick_endpoint(st); return build_request(p, k, "mixed-uniq")
    else
      local j = next_A(st); st.a_hits = st.a_hits + 1
      local p = pick_endpoint(st); return build_request(p, key_A(j), "mixed-A")
    end

  else -- uniform
    local k = next_unique(st)
    local p = pick_endpoint(st); return build_request(p, k, "uniform")
  end
end

request = function()
  local st = threads[wrk.thread]
  if not st then setup(wrk.thread); st = threads[wrk.thread] end
  local ok, req = pcall(generate_request, st)
  if ok and req then return req end
  -- Fallback на случай любой ошибки: всегда валидный запрос
  return wrk.format("GET",
    "/api/v2/pagedata" .. build_query(1),
    { ["Accept-Encoding"]=AE_HEADER, ["X-LFU-Phase"]="fallback" })
end

function done(summary, latency, requests)
  local agg={req=0,a_hits=0,b_hits=0,a_chal=0,b_chal=0,repeats=0}
  for _,st in pairs(threads) do
    agg.req=agg.req+st.req_total
    agg.a_hits=agg.a_hits+st.a_hits
    agg.b_hits=agg.b_hits+st.b_hits
    agg.a_chal=agg.a_chal+st.a_chal
    agg.b_chal=agg.b_chal+st.b_chal
    agg.repeats=agg.repeats+st.repeats
  end

  io.write("\n=== tiny-lfu (infinite) summary ===\n")
  io.write(string.format("mode=%s phase_mode=%s ring_pairs=%d\n", MODE, PHASE_MODE, HOT_RING))
  if PHASE_MODE=="time" then
    io.write(string.format("phases(time): warm=%ds battle=%ds shift=%ds\n", WARM_SEC,BATTLE_SEC,SHIFT_SEC))
  else
    io.write(string.format("phases(ops): warm=%d battle=%d shift=%d\n", OPS_WARM,OPS_BATTLE,OPS_SHIFT))
  end
  io.write(string.format("paths: mode=%s endpoints=%d AE=%s\n", PATH_MODE, #ENDPOINTS, AE_HEADER))
  io.write(string.format("requests: total=%d  A_hits=%d  B_hits=%d  A_chal=%d  B_chal=%d  repeats=%d\n",
    agg.req, agg.a_hits, agg.b_hits, agg.a_chal, agg.b_chal, agg.repeats))
  io.write(string.format("latency: mean=%.2fms stdev=%.2fms 99%%=%.2fms max=%.2fms\n",
    latency.mean/1000, latency.stdev/1000, latency:percentile(99.0)/1000, latency.max/1000))
  io.write(string.format("throughput: req/s=%.0f bytes/s=%.0f\n",
    summary.requests/(summary.duration/1000000), summary.bytes/(summary.duration/1000000)))
  io.write(string.format("errors: status=%d connect=%d read=%d write=%d timeout=%d\n\n",
    summary.errors.status, summary.errors.connect, summary.errors.read, summary.errors.write, summary.errors.timeout))

  io.write([[
EXPECT (server):
  • MODE=lfu:
      warm   → adm_allowed ≫ not_allowed; растёт entries; высокий hit_rate на A.
      battle → DK режет первые касания B; с двойным касанием B редко проходит LFU против A.
      shift  → B становится горячим; LFU адаптируется, A начинает проигрывать.
  • MODE=avalanche:
      бесконечный поток новых ключей → adm_not_allowed ≫ adm_allowed (DK+LFU фильтруют холод).
]])
end
