// Placement-path atomic scripts; see twins.ts for the Lua/JS twin contract.

import { sha1Hex } from "../hash.js";
import {
  HELPERS_LUA,
  liveOwner,
  parseList,
  parseObj,
  type Script,
} from "./twins.js";

/** New placements need this much lease left, matching pinned's worker-side
 * healthy buffer; capped at leaseTtl/3 so short-lease fleets stay placeable. */
export const PLACEMENT_HEALTH_MARGIN_MS = 5_000;

/** Fixed rebase point of the placement-epoch time arm (2026-01-01 UTC). */
export const PROJECT_EPOCH_MS = Date.UTC(2026, 0, 1);

const MINT_EPOCH_LUA = `
local function mintEpoch(floorKey, prevEpoch, now)
  local prev = tonumber(prevEpoch) or 0
  local floorRaw = redis.call('GET', floorKey)
  local floor = 0
  if floorRaw then
    floor = tonumber(floorRaw)
    if floor == nil or floor ~= floor or floor == math.huge or floor == -math.huge then
      error(redis.error_reply('malformed epoch_floor'))
    end
  end
  local epoch = math.max((now - ${PROJECT_EPOCH_MS}) * 1000, prev + 1, floor + 1)
  local encoded = string.format('%.0f', epoch)
  redis.call('SET', floorKey, encoded)
  return encoded
end
`;

const mintEpoch = (
  kv: Script.SyncKV,
  floorKey: string,
  prevEpoch: unknown,
  now: number,
): string => {
  const prev = typeof prevEpoch === "string" ? Number(prevEpoch) || 0 : 0;
  const floorRaw = kv.get(floorKey);
  const floor =
    floorRaw === null ? 0 : floorRaw.trim() === "" ? NaN : Number(floorRaw);
  if (!Number.isFinite(floor)) throw new Error("malformed epoch_floor");
  const epoch = Math.max((now - PROJECT_EPOCH_MS) * 1000, prev + 1, floor + 1);
  const encoded = String(epoch);
  kv.set(floorKey, encoded);
  return encoded;
};

/**
 * KEYS: placements, workers, ns_workers, boot_time, overwritten, epoch_floor
 * ARGV: id, now(ms), rand1, rand2, leaseTtlMs, ns,
 *       preferredSessionId(''=none), affinityKey(''=none), balanceFactor,
 *       hostname, env, takeover, excludeSessionId(''=none),
 *       ejectedSessionIdsJson(''=none; skipped while another candidate remains)
 * Returns [status, workerRecordRaw|'', aux, carriedExpiry|'', epoch|'',
 * fresh]: aux is the remaining quarantine ms on 'quarantine', else '1'/'0'
 * for the id sitting on the overwritten list; carriedExpiry is set only when
 * this call placed with a carried wind-down; epoch is the placement row's
 * epoch on 'forward'; fresh is '1' when this call created the placement row.
 * On 'takeover-pending' aux is '1' when this call set the fence and the 4th
 * slot carries the fence expiry.
 * A re-claim (evicting a prior row before placing anew) writes the id's
 * overwritten marker atomically, stamping the minted epoch onto it (backoff
 * meta survives a re-list); a dead row that cannot be re-claimed
 * yet (quarantine, no backends) is left in place as the durability marker.
 * Each new placement decrements the chosen worker's stored headroom, so a
 * concurrent burst derates it until the next heartbeat refresh.
 */
export const resolveScript: Script.Def = {
  lua: `${HELPERS_LUA}${MINT_EPOCH_LUA}
local placements, workers, nsWorkers = KEYS[1], KEYS[2], KEYS[3]
local bootTime, overwritten, epochFloor = KEYS[4], KEYS[5], KEYS[6]
local id, now = ARGV[1], tonumber(ARGV[2])
local rand1, rand2 = tonumber(ARGV[3]), tonumber(ARGV[4])
local leaseTtl, ns = tonumber(ARGV[5]), ARGV[6]
local preferredSid, affinityKey = ARGV[7], ARGV[8]
local balanceFactor = tonumber(ARGV[9])
local hostname, env, takeover = ARGV[10], ARGV[11], ARGV[12]
local excludeSid = ARGV[13]
local ejected = {}
if ARGV[14] and ARGV[14] ~= '' then
  for _, sid in ipairs(cjson.decode(ARGV[14])) do ejected[sid] = true end
end
local margin = math.min(${PLACEMENT_HEALTH_MARGIN_MS}, leaseTtl / 3)
local listedField = env .. '\\n' .. hostname .. '\\n' .. ns .. '\\n' .. id

local function listed()
  return redis.call('HEXISTS', overwritten, listedField) == 1 and '1' or '0'
end

local carry = nil
local prevEpoch = nil
local row = redis.call('HGET', placements, id)
if row then
  local ok, placed = pcall(cjson.decode, row)
  local raw, rec = nil, nil
  if ok and type(placed) == 'table' then
    raw, rec = liveOwner(workers, placed.sessionId, now)
    if type(placed.epoch) == 'string' then prevEpoch = placed.epoch end
  end
  local pending = ok and type(placed) == 'table'
    and type(placed.expiry) == 'number' and now < placed.expiry
  if raw and rec.hostname ~= hostname then
    if takeover == 'deny' then
      return { 'takeover-denied', '', '', '', '', '' }
    end
    if placed.expiry == nil then
      local fence = now + leaseTtl
      placed.expiry = fence
      redis.call('HSET', placements, id, cjson.encode(placed))
      return { 'takeover-pending', raw, '1', tostring(fence), '', '' }
    end
    if pending then
      return { 'takeover-pending', raw, '0', tostring(placed.expiry), '', '' }
    end
  elseif raw and (placed.expiry == nil or pending) then
    return { 'forward', raw, listed(), '', prevEpoch or '', '0' }
  end
  if not raw and pending then carry = placed.expiry end
end

local list = redis.call('HGET', nsWorkers, hostname .. ns)
if not list then
  return { 'unknown-ns', '', '', '', '', '' }
end

redis.call('SET', bootTime, ARGV[2], 'NX')
local boot = tonumber(redis.call('GET', bootTime))
if boot == nil then
  return redis.error_reply('malformed boot_time')
end
local remaining = boot + leaseTtl - now
if remaining > 0 then
  return { 'quarantine', '', tostring(remaining), '', '', '' }
end

local candidates, ejectedOnly = {}, {}
local ok, sessionIds = pcall(cjson.decode, list)
if ok and type(sessionIds) == 'table' then
  for _, sid in ipairs(sessionIds) do
    local raw, rec = liveOwner(workers, sid, now)
    if raw and not rec.draining and not rec.degradedSince and sid ~= excludeSid
      and rec.expiry - now >= margin then
      local c = { raw = raw, sessionId = sid, headroom = rec.headroom }
      if ejected[sid] then ejectedOnly[#ejectedOnly + 1] = c
      else candidates[#candidates + 1] = c end
    end
  end
end
if #candidates == 0 then candidates = ejectedOnly end
if #candidates == 0 then
  return { 'no-backends', '', listed(), '', '', '' }
end
if preferredSid == '' and affinityKey ~= '' then
  local best = ''
  for _, c in ipairs(candidates) do
    local digest = redis.sha1hex(affinityKey .. c.sessionId)
    if digest > best then best, preferredSid = digest, c.sessionId end
  end
end
local chosen = nil
if preferredSid ~= '' then
  local sum = 0
  for _, c in ipairs(candidates) do sum = sum + c.headroom end
  for _, c in ipairs(candidates) do
    if c.sessionId == preferredSid
      and c.headroom >= sum / #candidates / balanceFactor then
      chosen = c
    end
  end
end
if not chosen then
  chosen = candidates[1]
  if #candidates > 1 then
    local a = math.floor(rand1 * #candidates) + 1
    local b = math.floor(rand2 * (#candidates - 1)) + 1
    if b >= a then b = b + 1 end
    local first, second = candidates[a], candidates[b]
    chosen = first.headroom >= second.headroom and first or second
  end
end
local epoch = mintEpoch(epochFloor, prevEpoch, now)
if row then
  redis.call('HDEL', placements, id)
  local markerOk, marker = pcall(cjson.decode,
    redis.call('HGET', overwritten, listedField) or '{}')
  if not markerOk or type(marker) ~= 'table' then marker = {} end
  marker.epoch = epoch
  redis.call('HSET', overwritten, listedField, cjson.encode(marker))
end
local derated = cjson.decode(chosen.raw)
derated.headroom = derated.headroom - 1
redis.call('HSET', workers, chosen.sessionId, cjson.encode(derated))
redis.call('HSET', placements, id,
  cjson.encode({ sessionId = chosen.sessionId, expiry = carry, epoch = epoch }))
return { 'forward', chosen.raw, listed(),
  carry and tostring(carry) or '', epoch, '1' }
`,
  js: (kv, keys, argv) => {
    const [placementsKey, workersKey, nsWorkersKey] = [
      keys[0]!,
      keys[1]!,
      keys[2]!,
    ];
    const [bootTimeKey, overwrittenKey, epochFloorKey] = [
      keys[3]!,
      keys[4]!,
      keys[5]!,
    ];
    const [id, nowRaw, rand1Raw, rand2Raw, leaseTtlRaw, ns] = [
      argv[0]!,
      argv[1]!,
      argv[2]!,
      argv[3]!,
      argv[4]!,
      argv[5]!,
    ];
    let preferredSid = argv[6]!;
    const affinityKey = argv[7]!;
    const balanceFactor = Number(argv[8]!);
    const [hostname, env, takeover] = [argv[9]!, argv[10]!, argv[11]!];
    const excludeSid = argv[12]!;
    const ejectedRaw = argv[13] ?? "";
    const ejected = new Set<string>(
      ejectedRaw === "" ? [] : (JSON.parse(ejectedRaw) as string[]),
    );
    const now = Number(nowRaw);
    const margin = Math.min(
      PLACEMENT_HEALTH_MARGIN_MS,
      Number(leaseTtlRaw) / 3,
    );
    const listedField = `${env}\n${hostname}\n${ns}\n${id}`;
    const listed = (): string =>
      kv.hget(overwrittenKey, listedField) !== null ? "1" : "0";

    let carry: number | null = null;
    let prevEpoch: string | undefined;
    const row = kv.hget(placementsKey, id);
    if (row !== null) {
      const placed = parseObj(row);
      const owner =
        placed === null
          ? null
          : liveOwner(kv, workersKey, placed["sessionId"], now);
      if (typeof placed?.["epoch"] === "string") {
        prevEpoch = placed["epoch"];
      }
      const expiry = placed?.["expiry"];
      const pending = typeof expiry === "number" && now < expiry;
      if (owner !== null && owner.rec["hostname"] !== hostname) {
        if (takeover === "deny") return ["takeover-denied", "", "", "", "", ""];
        if (expiry === undefined) {
          const fence = now + Number(leaseTtlRaw);
          placed!["expiry"] = fence;
          kv.hset(placementsKey, id, JSON.stringify(placed));
          return ["takeover-pending", owner.raw, "1", String(fence), "", ""];
        }
        if (pending) {
          return ["takeover-pending", owner.raw, "0", String(expiry), "", ""];
        }
      } else if (owner !== null && (expiry === undefined || pending)) {
        return ["forward", owner.raw, listed(), "", prevEpoch ?? "", "0"];
      }
      if (owner === null && pending) carry = expiry as number;
    }

    const list = kv.hget(nsWorkersKey, `${hostname}${ns}`);
    if (list === null) return ["unknown-ns", "", "", "", "", ""];

    if (kv.get(bootTimeKey) === null) kv.set(bootTimeKey, nowRaw);
    const boot = Number(kv.get(bootTimeKey));
    if (!Number.isFinite(boot)) throw new Error("malformed boot_time");
    const remaining = boot + Number(leaseTtlRaw) - now;
    if (remaining > 0) return ["quarantine", "", String(remaining), "", "", ""];

    let candidates: { raw: string; sessionId: string; headroom: number }[] = [];
    const ejectedOnly: typeof candidates = [];
    const sessionIds = parseList(list);
    if (sessionIds !== null) {
      for (const sid of sessionIds) {
        if (sid === excludeSid) continue;
        const owner = liveOwner(kv, workersKey, sid, now);
        if (
          owner !== null &&
          owner.rec["draining"] !== true &&
          owner.rec["degradedSince"] === undefined &&
          (owner.rec["expiry"] as number) - now >= margin
        ) {
          (ejected.has(sid as string) ? ejectedOnly : candidates).push({
            raw: owner.raw,
            sessionId: sid as string,
            headroom: Number(owner.rec["headroom"]),
          });
        }
      }
    }
    if (candidates.length === 0) candidates = ejectedOnly;
    if (candidates.length === 0)
      return ["no-backends", "", listed(), "", "", ""];
    if (preferredSid === "" && affinityKey !== "") {
      let best = "";
      for (const c of candidates) {
        const digest = sha1Hex(affinityKey + c.sessionId);
        if (digest > best) [best, preferredSid] = [digest, c.sessionId];
      }
    }
    let chosen: (typeof candidates)[number] | null = null;
    if (preferredSid !== "") {
      const sum = candidates.reduce((acc, c) => acc + c.headroom, 0);
      for (const c of candidates) {
        if (
          c.sessionId === preferredSid &&
          c.headroom >= sum / candidates.length / balanceFactor
        ) {
          chosen = c;
        }
      }
    }
    if (chosen === null) {
      chosen = candidates[0]!;
      if (candidates.length > 1) {
        const a = Math.floor(Number(rand1Raw) * candidates.length);
        let b = Math.floor(Number(rand2Raw) * (candidates.length - 1));
        if (b >= a) b += 1;
        const first = candidates[a]!;
        const second = candidates[b]!;
        chosen = first.headroom >= second.headroom ? first : second;
      }
    }
    const epoch = mintEpoch(kv, epochFloorKey, prevEpoch, now);
    if (row !== null) {
      kv.hdel(placementsKey, id);
      const marker = parseObj(kv.hget(overwrittenKey, listedField)) ?? {};
      marker["epoch"] = epoch;
      kv.hset(overwrittenKey, listedField, JSON.stringify(marker));
    }
    const derated = parseObj(chosen.raw)!;
    derated["headroom"] = (derated["headroom"] as number) - 1;
    kv.hset(workersKey, chosen.sessionId, JSON.stringify(derated));
    kv.hset(
      placementsKey,
      id,
      JSON.stringify({
        sessionId: chosen.sessionId,
        ...(carry !== null && { expiry: carry }),
        epoch,
      }),
    );
    return [
      "forward",
      chosen.raw,
      listed(),
      carry === null ? "" : String(carry),
      epoch,
      "1",
    ];
  },
};

/**
 * KEYS: placements
 * ARGV: sessionId, prevSessionId(''=none), now(ms), id...
 * Register-time reclaim, restricted to process continuity: only rows still
 * owned by prevSessionId re-bind to sessionId, keeping the row epoch (same
 * holder across a lease blip — no writer transition, no mint). A pending
 * wind-down expiry carries onto the re-bound row. Every other id — absent,
 * expired, or owned by any other session — is left untouched and reported;
 * such rows re-place through the resolve claim path when traffic arrives.
 * Returns [claimedCount, [lost id...], [carried id, expiry, ...]].
 */
export const reclaimScript: Script.Def = {
  lua: `
local placements = KEYS[1]
local sessionId, prevSessionId = ARGV[1], ARGV[2]
local now = tonumber(ARGV[3])
local claimed, lost, carries = 0, {}, {}
for i = 4, #ARGV do
  local id = ARGV[i]
  local held = nil
  if prevSessionId ~= '' then
    local row = redis.call('HGET', placements, id)
    if row then
      local ok, placed = pcall(cjson.decode, row)
      if ok and type(placed) == 'table' and placed.sessionId == prevSessionId then
        held = placed
      end
    end
  end
  if held then
    claimed = claimed + 1
    local carry = nil
    if type(held.expiry) == 'number' and now < held.expiry then
      carry = held.expiry
    end
    held.sessionId = sessionId
    held.expiry = carry
    redis.call('HSET', placements, id, cjson.encode(held))
    if carry then
      carries[#carries + 1] = id
      carries[#carries + 1] = tostring(carry)
    end
  else
    lost[#lost + 1] = id
  end
end
return { claimed, lost, carries }
`,
  js: (kv, keys, argv) => {
    const placementsKey = keys[0]!;
    const [sessionId, prevSessionId] = [argv[0]!, argv[1]!];
    const now = Number(argv[2]!);
    let claimed = 0;
    const lost: string[] = [];
    const carries: string[] = [];
    for (const id of argv.slice(3)) {
      let held: Record<string, unknown> | null = null;
      if (prevSessionId !== "") {
        const placed = parseObj(kv.hget(placementsKey, id));
        if (placed !== null && placed["sessionId"] === prevSessionId) {
          held = placed;
        }
      }
      if (held !== null) {
        claimed += 1;
        const expiry = held["expiry"];
        const carry =
          typeof expiry === "number" && now < expiry ? expiry : null;
        held["sessionId"] = sessionId;
        if (carry === null) delete held["expiry"];
        else held["expiry"] = carry;
        kv.hset(placementsKey, id, JSON.stringify(held));
        if (carry !== null) carries.push(id, String(carry));
      } else {
        lost.push(id);
      }
    }
    return [claimed, lost, carries];
  },
};

/**
 * KEYS: placements
 * ARGV: deadline(epoch ms, '' clears), id...
 * Sets or clears the wind-down expiry on each placed id's row.
 * Returns [id, owning sessionId, ...] pairs; unplaced ids are skipped.
 */
export const windDownScript: Script.Def = {
  lua: `
local out = {}
for i = 2, #ARGV do
  local raw = redis.call('HGET', KEYS[1], ARGV[i])
  if raw then
    local ok, row = pcall(cjson.decode, raw)
    if ok and type(row) == 'table' and type(row.sessionId) == 'string' then
      if ARGV[1] == '' then row.expiry = nil
      else row.expiry = tonumber(ARGV[1]) end
      redis.call('HSET', KEYS[1], ARGV[i], cjson.encode(row))
      out[#out + 1] = ARGV[i]
      out[#out + 1] = row.sessionId
    end
  end
end
return out
`,
  js: (kv, keys, argv) => {
    const placementsKey = keys[0]!;
    const deadline = argv[0]!;
    const out: string[] = [];
    for (const id of argv.slice(1)) {
      const row = parseObj(kv.hget(placementsKey, id));
      if (row === null || typeof row["sessionId"] !== "string") continue;
      if (deadline === "") delete row["expiry"];
      else row["expiry"] = Number(deadline);
      kv.hset(placementsKey, id, JSON.stringify(row));
      out.push(id, row["sessionId"]);
    }
    return out;
  },
};

/**
 * KEYS: placements
 * ARGV: sessionId, (id, epoch)...
 * Compare-and-delete on (sessionId, epoch): HDELs each id only while its
 * row's sessionId matches and its epoch equals the quoted one ('' quotes an
 * epoch-less row). A mismatch leaves the row.
 * Returns the number of rows deleted.
 */
export const releaseScript: Script.Def = {
  lua: `
local released = 0
for i = 2, #ARGV, 2 do
  local id, epoch = ARGV[i], ARGV[i + 1]
  local raw = redis.call('HGET', KEYS[1], id)
  if raw then
    local ok, row = pcall(cjson.decode, raw)
    if ok and type(row) == 'table' and row.sessionId == ARGV[1] then
      local matches
      if epoch == '' then
        matches = row.epoch == nil
      else
        matches = tonumber(epoch) ~= nil and type(row.epoch) == 'string'
          and tonumber(row.epoch) == tonumber(epoch)
      end
      if matches then
        released = released + redis.call('HDEL', KEYS[1], id)
      end
    end
  end
end
return released
`,
  js: (kv, keys, argv) => {
    const placementsKey = keys[0]!;
    const sessionId = argv[0]!;
    let released = 0;
    for (let i = 1; i < argv.length; i += 2) {
      const [id, epoch] = [argv[i]!, argv[i + 1]!];
      const row = parseObj(kv.hget(placementsKey, id));
      if (row === null || row["sessionId"] !== sessionId) continue;
      const rowEpoch =
        typeof row["epoch"] === "string" ? Number(row["epoch"]) : NaN;
      const matches =
        epoch === ""
          ? row["epoch"] === undefined
          : Number.isFinite(Number(epoch)) && rowEpoch === Number(epoch);
      if (matches) {
        kv.hdel(placementsKey, id);
        released += 1;
      }
    }
    return released;
  },
};

/**
 * KEYS: overwritten
 * ARGV: field, epoch, metaJson('' = delete)
 * Compare-and-write fenced on the marker's stamped epoch: '' HDELs the field
 * while its epoch <= the quoted one (a proof of life under epoch E clears
 * every generation listed at or before E); a metaJson re-writes the backoff
 * meta only while the stored epoch equals the quoted one. A marker without a
 * numeric epoch is an error. Returns 1 on write, 0 on skip.
 */
export const overwrittenCasScript: Script.Def = {
  lua: `
local raw = redis.call('HGET', KEYS[1], ARGV[1])
if not raw then return 0 end
local ok, marker = pcall(cjson.decode, raw)
local stored = ok and type(marker) == 'table' and tonumber(marker.epoch) or nil
local quoted = tonumber(ARGV[2])
if not stored or not quoted then
  return redis.error_reply('malformed overwritten marker epoch')
end
if ARGV[3] == '' then
  if stored <= quoted then return redis.call('HDEL', KEYS[1], ARGV[1]) end
  return 0
end
if stored == quoted then
  redis.call('HSET', KEYS[1], ARGV[1], ARGV[3])
  return 1
end
return 0
`,
  js: (kv, keys, argv) => {
    const overwrittenKey = keys[0]!;
    const [field, epochRaw, metaJson] = [argv[0]!, argv[1]!, argv[2]!];
    const raw = kv.hget(overwrittenKey, field);
    if (raw === null) return 0;
    const markerEpoch = parseObj(raw)?.["epoch"];
    const stored =
      typeof markerEpoch === "string" && markerEpoch !== ""
        ? Number(markerEpoch)
        : typeof markerEpoch === "number"
          ? markerEpoch
          : NaN;
    const quoted = epochRaw === "" ? NaN : Number(epochRaw);
    if (!Number.isFinite(stored) || !Number.isFinite(quoted)) {
      throw new Error("malformed overwritten marker epoch");
    }
    if (metaJson === "") {
      if (stored > quoted) return 0;
      kv.hdel(overwrittenKey, field);
      return 1;
    }
    if (stored !== quoted) return 0;
    kv.hset(overwrittenKey, field, metaJson);
    return 1;
  },
};