## Pure unit guard for the melee/1v1 rack-membership mechanism.
##
## No Java, no battle, no virtual bullets: the fitness windows are seeded
## directly so the selection is deterministic and fast. Covers:
##   * mode filtering selects the right gun subset,
##   * an EMPTY filtered set falls back to the full rack,
##   * `TR_RACK_<GUN>` env overrides parse correctly (both / 1v1 / melee / off),
##   * the shipped default is `onlyPattern` (Pattern admitted, all else `off`),
##   * the revert override restores the old full rack,
##   * an all-`both` membership reproduces the pre-change selection exactly.
##
## Run with plain:
##   nim c -r common_libs/tests/test_rack_membership.nim

import std/[random, tables, os, strutils]
import gun_harness/gun_interface
import gun_harness/virtual_bullets
import gun_harness/selector

var failures = 0
proc check(name: string, ok: bool) =
  if ok: echo "PASS: ", name
  else: echo "FAIL: ", name; inc failures

proc recordHit(fw: var FitnessWindow, hit: bool) =
  fw.hits[fw.head] = hit
  fw.head = (fw.head + 1) mod WindowSize
  inc fw.count

proc seedWindow(t: var VirtualTracker, targetId, gunId, binIdx, hits, misses: int) =
  if targetId notin t.fitness:
    t.fitness[targetId] = newSeq[GunFitness](t.numGuns)
  var fw = addr t.fitness[targetId][gunId].bins[binIdx]
  for _ in 0..<hits:   recordHit(fw[], true)
  for _ in 0..<misses: recordHit(fw[], false)

# ── mode derivation (server truth) ────────────────────────────────────────────

proc testModeFromEnemyCount() =
  check "rackMode: 1 enemy alive -> 1v1",
        rackMode(1) == rm1v1
  check "rackMode: 2+ enemies alive -> melee",
        rackMode(2) == rmMelee and rackMode(5) == rmMelee
  check "rackMode: 0 enemies (should not occur) stays melee, matching the radar",
        rackMode(0) == rmMelee

# ── pure filter ────────────────────────────────────────────────────────────────

proc testAdmittedGunsFiltering() =
  # 0: both, 1: only1v1, 2: onlyMelee, 3: off
  let mem = [rmBoth, rmOnly1v1, rmOnlyMelee, rmOff]
  check "admittedGuns(1v1) keeps both + only1v1, drops onlyMelee/off",
        admittedGuns(4, rm1v1, mem) == @[0, 1]
  check "admittedGuns(melee) keeps both + onlyMelee, drops only1v1/off",
        admittedGuns(4, rmMelee, mem) == @[0, 2]

proc testEmptyFilterFallsBackToFullRack() =
  let allOff = [rmOff, rmOff, rmOff, rmOff]
  check "admittedGuns: all-off rack degrades to the FULL rack (1v1)",
        admittedGuns(4, rm1v1, allOff) == @[0, 1, 2, 3]
  check "admittedGuns: all-off rack degrades to the FULL rack (melee)",
        admittedGuns(4, rmMelee, allOff) == @[0, 1, 2, 3]
  # A rack where only melee guns exist still falls back to all guns in 1v1.
  let onlyMelee = [rmOnlyMelee, rmOnlyMelee]
  check "admittedGuns: a rack with NO gun for the mode degrades to the full rack",
        admittedGuns(2, rm1v1, onlyMelee) == @[0, 1]

proc testEmptyMembershipMeansNoFiltering() =
  check "admittedGuns: an empty membership table filters nothing",
        admittedGuns(3, rm1v1, []) == @[0, 1, 2]
  check "rackAdmitted: an empty membership table admits every id",
        rackAdmitted(0, rm1v1, []) and rackAdmitted(99, rmMelee, [])

proc testRackAdmitted() =
  let mem = [rmBoth, rmOnly1v1, rmOff]
  check "rackAdmitted: both is admitted in every mode",
        rackAdmitted(0, rm1v1, mem) and rackAdmitted(0, rmMelee, mem)
  check "rackAdmitted: only1v1 is admitted only in 1v1",
        rackAdmitted(1, rm1v1, mem) and not rackAdmitted(1, rmMelee, mem)
  check "rackAdmitted: off is never admitted",
        not rackAdmitted(2, rm1v1, mem) and not rackAdmitted(2, rmMelee, mem)
  check "rackAdmitted: negative id is never admitted",
        not rackAdmitted(-1, rm1v1, mem)

# ── selection-level filtering ─────────────────────────────────────────────────

proc testSelectionFollowsMode() =
  # gun0 50%, gun1 100% (1v1 only), gun2 100% (melee only).
  var t = initTracker(3)
  seedWindow(t, 7, 0, 0, 25, 25)
  seedWindow(t, 7, 1, 0, 50, 0)
  seedWindow(t, 7, 2, 0, 50, 0)
  let mem = @[rmBoth, rmOnly1v1, rmOnlyMelee]
  randomize(7)
  var one: array[3, int]
  for _ in 0..<100: inc one[t.bestGun(7, rackMode = rm1v1, membership = mem)]
  var melee: array[3, int]
  for _ in 0..<100: inc melee[t.bestGun(7, rackMode = rmMelee, membership = mem)]
  check "selection: 1v1 rack only ever returns admitted guns {0,1}",
        one[2] == 0 and one[0] + one[1] == 100
  check "selection: 1v1 rack picks the only1v1 gun when it dominates",
        one[1] == 100
  check "selection: melee rack only ever returns admitted guns {0,2}",
        melee[1] == 0 and melee[0] + melee[2] == 100
  check "selection: melee rack picks the onlyMelee gun when it dominates",
        melee[2] == 100

proc testSelectionEmptyFilterFallsBack() =
  var t = initTracker(3)
  seedWindow(t, 7, 0, 0, 10, 40)  # 20%
  seedWindow(t, 7, 1, 0, 50, 0)   # 100%
  seedWindow(t, 7, 2, 0, 10, 40)  # 20%
  let allOff = @[rmOff, rmOff, rmOff]
  randomize(11)
  var seen: array[3, int]
  for _ in 0..<100: inc seen[t.bestGun(7, rackMode = rm1v1, membership = allOff)]
  check "selection: an all-off rack still selects (falls back to full rack)",
        seen[0] + seen[1] + seen[2] == 100
  check "selection: the fallback still picks the best gun (gun 1)",
        seen[1] == 100

# ── env parsing / defaults ────────────────────────────────────────────────────

proc clearRackEnv() =
  for name in RackGunNames:
    delEnv("TR_RACK_" & name)

proc testParseMembership() =
  check "parse: 'both' -> rmBoth", parseRackMembership("both") == rmBoth
  check "parse: '1v1' -> rmOnly1v1", parseRackMembership("1v1") == rmOnly1v1
  check "parse: 'melee' -> rmOnlyMelee", parseRackMembership("melee") == rmOnlyMelee
  check "parse: 'off' -> rmOff", parseRackMembership("off") == rmOff
  check "parse: case/whitespace-insensitive",
        parseRackMembership("  1V1 ") == rmOnly1v1 and
        parseRackMembership("MeLeE") == rmOnlyMelee
  check "parse: empty -> shipped default rmBoth",
        parseRackMembership("") == rmBoth
  check "parse: unknown -> rmBoth (never fatal)",
        parseRackMembership("banana") == rmBoth

proc testDefaultsOnlyPattern() =
  clearRackEnv()
  # SHIPPED DEFAULT: `onlyPattern`. Pattern (id 5) is `both`; every other gun is
  # `off`. The selector mechanism is untouched, but the rack has one member.
  const PatternId = 5
  var exactlyOne = true
  var admittedCount = 0
  for i in 0..<RackGunNames.len:
    if DefaultRackMembership[i] != rmOff:
      inc admittedCount
      if i != PatternId or DefaultRackMembership[i] != rmBoth: exactlyOne = false
  check "defaults: the shipped table admits exactly one gun", admittedCount == 1
  check "defaults: that one gun is PATTERN (id 5), in both racks",
        exactlyOne and DefaultRackMembership[PatternId] == rmBoth
  check "defaults: every non-Pattern gun is `off`",
        DefaultRackMembership[0] == rmOff and DefaultRackMembership[1] == rmOff and
        DefaultRackMembership[4] == rmOff and DefaultRackMembership[6] == rmOff and
        DefaultRackMembership[12] == rmOff and DefaultRackMembership[13] == rmOff and
        DefaultRackMembership[14] == rmOff
  check "defaults: the default rack admits ONLY Pattern in 1v1",
        admittedGuns(18, rm1v1, DefaultRackMembership) == @[PatternId]
  check "defaults: the default rack admits ONLY Pattern in melee",
        admittedGuns(18, rmMelee, DefaultRackMembership) == @[PatternId]
  let loaded = loadRackMembership()
  check "defaults: with a clean environment loadRackMembership() == shipped table",
        loaded == DefaultRackMembership
  check "defaults: RackGunNames covers the shipped 18-gun rack",
        RackGunNames.len == 18 and DefaultRackMembership.len == 18

proc testFloorRespectsAdmission() =
  ## The FLOOR path (`bestRate <= 0` or below the floor rate) historically fell
  ## back to gun 0 (HeadOn). `chooseFromFit` returns `admitted[0]`, so with the
  ## onlyPattern default (HeadOn `off`) the floor must return Pattern, an
  ## ADMITTED gun — never the removed HeadOn.
  var t = initTracker(15)
  # HeadOn strong but `off`; every admitted gun cold, so bestRate == 0 -> floor.
  seedWindow(t, 7, 0, 0, 50, 0)
  var diag: SelectorDiag
  let g = t.bestGun(7, addr diag, rackMode = rm1v1,
                    membership = DefaultRackMembership)
  check "floor: fires when the only admitted gun has no data", diag.floorFired
  check "floor: returns the ADMITTED gun (Pattern, id 5), NOT HeadOn", g == 5
  # Same field, same floor, but with HeadOn admissible: it returns gun 0, which
  # is exactly WHY HeadOn had to be removed for the floor to be safe.
  var oldAll: array[15, RackMembership]
  for i in 0..<15: oldAll[i] = rmBoth
  var t2 = initTracker(15)
  seedWindow(t2, 7, 0, 0, 0, 50)  # HeadOn 0%, eligible (>= MinObs)
  var d2: SelectorDiag
  let gOld = t2.bestGun(7, addr d2, rackMode = rm1v1, membership = oldAll)
  check "floor: with HeadOn admitted the old floor returns gun 0",
        d2.floorFired and gOld == 0
  # The LIVE path (`selectGun`) takes the same floor branch.
  var t3 = initTracker(15)
  seedWindow(t3, 7, 0, 0, 50, 0)
  var d3: SelectorDiag
  let gLive = t3.selectGun(7, tick = 0, diag = addr d3, rackMode = rm1v1,
                           membership = DefaultRackMembership)
  check "floor: the live selectGun floor also returns Pattern",
        d3.floorFired and gLive == 5

proc testRevertOverrideRestoresFullRack() =
  ## The documented one-liner: every TR_RACK_<GUN> back to `both` (TMPATTERN
  ## stays `off`). This must reproduce the old full rack without a rebuild.
  clearRackEnv()
  putEnv("TR_RACK_PATTERN", "both")
  # The ADE+SBC gun (id 17) additionally needs TR_BITBRAIN_NET=1; without the
  # switch TR_RACK_BITBRAIN is the LEGACY name for LEADGAIN (id 16), so the
  # documented one-liner has to flip the namespace over as well.
  putEnv("TR_BITBRAIN_NET", "1")
  for name in RackGunNames:
    if name != "TMPATTERN": putEnv("TR_RACK_" & name, "both")
  let m = loadRackMembership()
  delEnv("TR_BITBRAIN_NET")
  var full = true
  for i in 0..<RackGunNames.len:
    let want = if i == 14: rmOff else: rmBoth
    if m[i] != want: full = false
  check "revert: the documented one-liner restores the all-`both` full rack", full
  check "revert: 1v1 rack admits every gun again (TMPATTERN excluded)",
        admittedGuns(18, rm1v1, m) == @[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 15, 16, 17]
  check "revert: the legacy TR_RACK_BITBRAIN alone still selects LEADGAIN only",
        loadRackMembership()[17] == rmOff
  clearRackEnv()

proc testEnvOverrides() =
  clearRackEnv()
  putEnv("TR_RACK_TSETLIN", "1v1")
  putEnv("TR_RACK_DISPLACE", "melee")
  putEnv("TR_RACK_HEADON", "off")
  let m = loadRackMembership()
  check "env: TR_RACK_TSETLIN=1v1 moves only Tsetlin (id 2)",
        m[2] == rmOnly1v1
  check "env: TR_RACK_DISPLACE=melee moves only Displace (id 9)",
        m[9] == rmOnlyMelee
  check "env: TR_RACK_HEADON=off removes only HeadOn (id 0)",
        m[0] == rmOff
  check "env: every untouched gun keeps its shipped default (off)",
        m[1] == rmOff and m[3] == rmOff and m[13] == rmOff and m[14] == rmOff and
        m[5] == rmBoth
  putEnv("TR_RACK_KNN", "banana")
  check "env: unknown values are ignored, not fatal",
        loadRackMembership()[12] == rmBoth
  clearRackEnv()

# ── parity: default path reproduces pre-change selection ─────────────────────

proc testDefaultReproducesPreChange() =
  ## The pre-change selector is `bestGun(t, id)` (empty membership). With the
  ## all-`both` default table and EITHER mode, the selection sequence must be
  ## identical, RNG draw for RNG draw.
  var t = initTracker(5)
  seedWindow(t, 7, 0, 0, 60, 40)
  seedWindow(t, 7, 1, 0, 40, 60)
  seedWindow(t, 7, 2, 0, 80, 20)
  seedWindow(t, 7, 3, 0, 55, 45)
  seedWindow(t, 7, 4, 0, 20, 80)
  let allBoth = @[rmBoth, rmBoth, rmBoth, rmBoth, rmBoth]
  randomize(1234)
  var pre: seq[int]
  for _ in 0..<300: pre.add t.bestGun(7)
  randomize(1234)
  var as1v1: seq[int]
  for _ in 0..<300: as1v1.add t.bestGun(7, rackMode = rm1v1, membership = allBoth)
  randomize(1234)
  var asMelee: seq[int]
  for _ in 0..<300: asMelee.add t.bestGun(7, rackMode = rmMelee, membership = allBoth)
  check "parity: all-`both` membership reproduces pre-change selection (1v1)",
        pre == as1v1
  check "parity: all-`both` membership reproduces pre-change selection (melee)",
        pre == asMelee

proc testDefaultReproducesPreChangeSelectGun() =
  ## Same parity check through the live hysteresis path (`selectGun`).
  var t = initTracker(4)
  seedWindow(t, 7, 0, 0, 60, 40)
  seedWindow(t, 7, 1, 0, 70, 30)
  seedWindow(t, 7, 2, 0, 40, 60)
  seedWindow(t, 7, 3, 0, 50, 50)
  let allBoth = @[rmBoth, rmBoth, rmBoth, rmBoth]
  randomize(99)
  var a = initTracker(4)
  a = t
  var pre: seq[int]
  for tick in 0..<300: pre.add a.selectGun(7, tick = tick)
  randomize(99)
  var b = initTracker(4)
  b = t
  var post: seq[int]
  for tick in 0..<300:
    post.add b.selectGun(7, tick = tick, rackMode = rmMelee, membership = allBoth)
  check "parity: selectGun with all-`both` membership is byte-identical",
        pre == post

# ── driver ────────────────────────────────────────────────────────────────────

testModeFromEnemyCount()
testAdmittedGunsFiltering()
testEmptyFilterFallsBackToFullRack()
testEmptyMembershipMeansNoFiltering()
testRackAdmitted()
testSelectionFollowsMode()
testSelectionEmptyFilterFallsBack()
testParseMembership()
testDefaultsOnlyPattern()
testFloorRespectsAdmission()
testRevertOverrideRestoresFullRack()
testEnvOverrides()
testDefaultReproducesPreChange()
testDefaultReproducesPreChangeSelectGun()

if failures > 0:
  echo "\n", failures, " check(s) FAILED"
  quit(1)
echo "\nAll rack-membership checks passed."