// Worker-directory scripts (workers, ns_workers, bindings); see twins.ts for
// the Lua/JS twin contract.

import {
  bindEpoch,
  drainBase,
  HELPERS_LUA,
  parseObj,
  parseList,
  pruneWorkers,
  rewriteNs,
  type Script,
} from "./twins.js";

/**
 * KEYS: workers, ns_workers, client config, bindings, bindepochs, client metadata
 * ARGV: recordJson (sans draining/expiry), now(ms), leaseTtlMs,
 *       selfDraining('1'/'0'), epoch
 * Bucket-0 register: prunes long-dead workers (GC'ing unreferenced bindings
 * and their epoch entries), check-or-sets the per-namespace (hostname, ns)/env
 * bindings — a higher epoch displaces a disagreeing binding, quarantining the
 * (hostname, ns) for one lease — computes draining/expiry against the drain
 * config, writes the record, and rewrites its ns lists.
 * Returns ['ok', recordRaw, epochEntriesJson] — mirrors write both verbatim —
 * or ['conflict', message] when a binding disagrees at an equal or higher epoch.
 */
export const registerScript: Script.Def = {
  lua: `${HELPERS_LUA}
local rec = cjson.decode(ARGV[1])
local now, lease = tonumber(ARGV[2]), tonumber(ARGV[3])
local epoch = tonumber(ARGV[5])
local existingMeta = parseTable(redis.call('GET', KEYS[6]))
if rec.account and existingMeta and existingMeta.account
  and existingMeta.account ~= rec.account then
  return { 'conflict', 'client belongs to another account' }
end
pruneWorkers(KEYS[1], KEYS[4], KEYS[5], now, lease)
for _, ns in ipairs(rec.namespaces) do
  local env = rec.envs[ns]
  local boundEnv = redis.call('HGET', KEYS[4], 'h:' .. rec.hostname .. ns)
  if boundEnv and boundEnv ~= env then
    local bound = bindEpoch(KEYS[5], 'h:' .. rec.hostname .. ns)
    if epoch <= bound.epoch then
      return { 'conflict', 'hostname ' .. rec.hostname .. ns ..
        ' is bound to env ' .. boundEnv .. ' at epoch ' .. bound.epoch ..
        ', not ' .. env .. ' — register with a higher epoch to displace' }
    end
  end
  local boundTakeover = redis.call('HGET', KEYS[4], 'e:' .. env)
  if boundTakeover and boundTakeover ~= rec.takeover then
    local bound = bindEpoch(KEYS[5], 'e:' .. env)
    if epoch <= bound.epoch then
      return { 'conflict', 'env ' .. env .. ' agreed takeover ' ..
        boundTakeover .. ' at epoch ' .. bound.epoch .. ', not ' ..
        rec.takeover .. ' — register with a higher epoch to displace' }
    end
  end
end
if rec.account then
  redis.call('SET', KEYS[6], cjson.encode({
    account = rec.account,
    hostname = rec.hostname,
    namespaces = rec.namespaces,
    envs = rec.envs,
  }))
end
local entries = {}
for _, ns in ipairs(rec.namespaces) do
  local env = rec.envs[ns]
  local hField = 'h:' .. rec.hostname .. ns
  local boundEnv = redis.call('HGET', KEYS[4], hField)
  local hBound = bindEpoch(KEYS[5], hField)
  local hEntry = { epoch = math.max(hBound.epoch, epoch) }
  if boundEnv and boundEnv ~= env then
    hEntry.quarantineUntil = now + lease
  elseif type(hBound.quarantineUntil) == 'number'
    and hBound.quarantineUntil > now then
    hEntry.quarantineUntil = hBound.quarantineUntil
  end
  redis.call('HSET', KEYS[4], hField, env)
  if hEntry.epoch > 0 or hEntry.quarantineUntil then
    entries[hField] = cjson.encode(hEntry)
    redis.call('HSET', KEYS[5], hField, entries[hField])
  end
  local eField = 'e:' .. env
  local eBound = bindEpoch(KEYS[5], eField)
  local eEpoch = math.max(eBound.epoch, epoch)
  redis.call('HSET', KEYS[4], eField, rec.takeover)
  if eEpoch > 0 then
    entries[eField] = cjson.encode({ epoch = eEpoch })
    redis.call('HSET', KEYS[5], eField, entries[eField])
  end
end
local draining, base = drainBase(redis.call('GET', KEYS[3]), now)
rec.draining = draining or ARGV[4] == '1'
rec.expiry = base + lease
local raw = cjson.encode(rec)
redis.call('HSET', KEYS[1], rec.sessionId, raw)
for _, ns in ipairs(rec.namespaces) do
  rewriteNs(KEYS[1], KEYS[2], rec.hostname .. ns, rec.sessionId, nil, now)
end
return { 'ok', raw, cjson.encode(entries) }
`,
  js: (kv, keys, argv) => {
    const [
      workersKey,
      nsWorkersKey,
      clientKey,
      bindingsKey,
      epochsKey,
      metaKey,
    ] = [keys[0]!, keys[1]!, keys[2]!, keys[3]!, keys[4]!, keys[5]!];
    const rec = JSON.parse(argv[0]!) as Record<string, unknown>;
    const now = Number(argv[1]!);
    const lease = Number(argv[2]!);
    const epoch = Number(argv[4]!);
    const [hostname, takeover] = [
      rec["hostname"] as string,
      rec["takeover"] as string,
    ];
    const envs = rec["envs"] as Record<string, string>;
    const existingMeta = parseObj(kv.get(metaKey));
    const account = rec["account"];
    if (
      typeof account === "string" &&
      typeof existingMeta?.["account"] === "string" &&
      existingMeta["account"] !== account
    ) {
      return ["conflict", "client belongs to another account"];
    }
    pruneWorkers(kv, workersKey, bindingsKey, epochsKey, now, lease);
    for (const ns of rec["namespaces"] as string[]) {
      const env = envs[ns]!;
      const boundEnv = kv.hget(bindingsKey, `h:${hostname}${ns}`);
      if (boundEnv !== null && boundEnv !== env) {
        const bound = bindEpoch(kv, epochsKey, `h:${hostname}${ns}`);
        if (epoch <= (bound["epoch"] as number)) {
          return [
            "conflict",
            `hostname ${hostname}${ns} is bound to env ${boundEnv} at epoch ` +
              `${String(bound["epoch"])}, not ${env} — register with a ` +
              `higher epoch to displace`,
          ];
        }
      }
      const boundTakeover = kv.hget(bindingsKey, `e:${env}`);
      if (boundTakeover !== null && boundTakeover !== takeover) {
        const bound = bindEpoch(kv, epochsKey, `e:${env}`);
        if (epoch <= (bound["epoch"] as number)) {
          return [
            "conflict",
            `env ${env} agreed takeover ${boundTakeover} at epoch ` +
              `${String(bound["epoch"])}, not ${takeover} — register with a ` +
              `higher epoch to displace`,
          ];
        }
      }
    }
    if (typeof account === "string") {
      kv.set(
        metaKey,
        JSON.stringify({
          account,
          hostname,
          namespaces: rec["namespaces"],
          envs,
        }),
      );
    }
    const entries: Record<string, string> = {};
    for (const ns of rec["namespaces"] as string[]) {
      const env = envs[ns]!;
      const hField = `h:${hostname}${ns}`;
      const boundEnv = kv.hget(bindingsKey, hField);
      const hBound = bindEpoch(kv, epochsKey, hField);
      const hEntry: Record<string, number> = {
        epoch: Math.max(hBound["epoch"] as number, epoch),
      };
      if (boundEnv !== null && boundEnv !== env) {
        hEntry["quarantineUntil"] = now + lease;
      } else if (
        typeof hBound["quarantineUntil"] === "number" &&
        hBound["quarantineUntil"] > now
      ) {
        hEntry["quarantineUntil"] = hBound["quarantineUntil"];
      }
      kv.hset(bindingsKey, hField, env);
      if (hEntry["epoch"]! > 0 || hEntry["quarantineUntil"] !== undefined) {
        entries[hField] = JSON.stringify(hEntry);
        kv.hset(epochsKey, hField, entries[hField]!);
      }
      const eField = `e:${env}`;
      const eBound = bindEpoch(kv, epochsKey, eField);
      const eEpoch = Math.max(eBound["epoch"] as number, epoch);
      kv.hset(bindingsKey, eField, takeover);
      if (eEpoch > 0) {
        entries[eField] = JSON.stringify({ epoch: eEpoch });
        kv.hset(epochsKey, eField, entries[eField]!);
      }
    }
    const drain = drainBase(kv.get(clientKey), now);
    rec["draining"] = drain.draining || argv[3] === "1";
    rec["expiry"] = drain.base + lease;
    const raw = JSON.stringify(rec);
    const sessionId = rec["sessionId"] as string;
    kv.hset(workersKey, sessionId, raw);
    for (const ns of rec["namespaces"] as string[]) {
      rewriteNs(
        kv,
        workersKey,
        nsWorkersKey,
        `${hostname}${ns}`,
        sessionId,
        null,
        now,
      );
    }
    return ["ok", raw, JSON.stringify(entries)];
  },
};

/**
 * KEYS: workers, ns_workers, bindings, bindepochs
 * ARGV: recordRaw (as returned by register), now(ms), leaseTtlMs,
 *       epochEntriesJson (as returned by register)
 * Buckets 1..S-1: writes the bucket-0 record, bindings, and binding-epoch
 * entries byte-for-byte, rewrites its ns lists, prunes long-dead workers.
 */
export const mirrorScript: Script.Def = {
  lua: `${HELPERS_LUA}
local rec = cjson.decode(ARGV[1])
local now, lease = tonumber(ARGV[2]), tonumber(ARGV[3])
redis.call('HSET', KEYS[1], rec.sessionId, ARGV[1])
for _, ns in ipairs(rec.namespaces) do
  redis.call('HSET', KEYS[3], 'h:' .. rec.hostname .. ns, rec.envs[ns])
  redis.call('HSET', KEYS[3], 'e:' .. rec.envs[ns], rec.takeover)
  rewriteNs(KEYS[1], KEYS[2], rec.hostname .. ns, rec.sessionId, nil, now)
end
for field, value in pairs(cjson.decode(ARGV[4])) do
  redis.call('HSET', KEYS[4], field, value)
end
pruneWorkers(KEYS[1], KEYS[3], KEYS[4], now, lease)
return 1
`,
  js: (kv, keys, argv) => {
    const [workersKey, nsWorkersKey, bindingsKey, epochsKey] = [
      keys[0]!,
      keys[1]!,
      keys[2]!,
      keys[3]!,
    ];
    const rec = JSON.parse(argv[0]!) as Record<string, unknown>;
    const now = Number(argv[1]!);
    const lease = Number(argv[2]!);
    const sessionId = rec["sessionId"] as string;
    const hostname = rec["hostname"] as string;
    const envs = rec["envs"] as Record<string, string>;
    kv.hset(workersKey, sessionId, argv[0]!);
    for (const ns of rec["namespaces"] as string[]) {
      kv.hset(bindingsKey, `h:${hostname}${ns}`, envs[ns]!);
      kv.hset(bindingsKey, `e:${envs[ns]!}`, rec["takeover"] as string);
      rewriteNs(
        kv,
        workersKey,
        nsWorkersKey,
        `${hostname}${ns}`,
        sessionId,
        null,
        now,
      );
    }
    for (const [field, value] of Object.entries(
      JSON.parse(argv[3]!) as Record<string, string>,
    )) {
      kv.hset(epochsKey, field, value);
    }
    pruneWorkers(kv, workersKey, bindingsKey, epochsKey, now, lease);
    return 1;
  },
};

/**
 * KEYS: workers, ns_workers, bindings, bindepochs
 * ARGV: recordRaw (refreshed record), now(ms), entriesJson (''=bucket 0)
 * Per-bucket heartbeat refresh: writes the record and ensures the session's
 * ns-list membership, so a bucket added by a shard-count raise serves without
 * re-register. Bucket 0 returns its {bindings, epochs} for the record's
 * fields; buckets 1..S-1 fill absent bindings from them (first writer wins —
 * a disagreeing binding and its epoch entry are left untouched; displacement
 * stays register-only).
 */
export const heartbeatScript: Script.Def = {
  lua: `${HELPERS_LUA}
local rec = cjson.decode(ARGV[1])
local now = tonumber(ARGV[2])
redis.call('HSET', KEYS[1], rec.sessionId, ARGV[1])
for _, ns in ipairs(rec.namespaces) do
  local field = rec.hostname .. ns
  local member = false
  local list = parseTable(redis.call('HGET', KEYS[2], field))
  if list then
    for _, sid in ipairs(list) do
      if sid == rec.sessionId then member = true end
    end
  end
  if not member then
    rewriteNs(KEYS[1], KEYS[2], field, rec.sessionId, nil, now)
  end
end
if ARGV[3] == '' then
  local out = { bindings = {}, epochs = {} }
  for _, ns in ipairs(rec.namespaces) do
    for _, field in ipairs({ 'h:' .. rec.hostname .. ns,
      'e:' .. rec.envs[ns] }) do
      local bound = redis.call('HGET', KEYS[3], field)
      if bound then out.bindings[field] = bound end
      local entry = redis.call('HGET', KEYS[4], field)
      if entry then out.epochs[field] = entry end
    end
  end
  return cjson.encode(out)
end
local entries = cjson.decode(ARGV[3])
for field, value in pairs(entries.bindings) do
  local cur = redis.call('HGET', KEYS[3], field)
  if not cur then
    redis.call('HSET', KEYS[3], field, value)
    cur = value
  end
  if cur == value and entries.epochs[field] then
    redis.call('HSETNX', KEYS[4], field, entries.epochs[field])
  end
end
return 1
`,
  js: (kv, keys, argv) => {
    const [workersKey, nsWorkersKey, bindingsKey, epochsKey] = [
      keys[0]!,
      keys[1]!,
      keys[2]!,
      keys[3]!,
    ];
    const rec = JSON.parse(argv[0]!) as Record<string, unknown>;
    const now = Number(argv[1]!);
    const sessionId = rec["sessionId"] as string;
    const hostname = rec["hostname"] as string;
    const envs = rec["envs"] as Record<string, string>;
    kv.hset(workersKey, sessionId, argv[0]!);
    for (const ns of rec["namespaces"] as string[]) {
      const field = `${hostname}${ns}`;
      const list = parseList(kv.hget(nsWorkersKey, field));
      if (list === null || !list.includes(sessionId)) {
        rewriteNs(kv, workersKey, nsWorkersKey, field, sessionId, null, now);
      }
    }
    if (argv[2] === "") {
      const out = {
        bindings: {} as Record<string, string>,
        epochs: {} as Record<string, string>,
      };
      for (const ns of rec["namespaces"] as string[]) {
        for (const field of [`h:${hostname}${ns}`, `e:${envs[ns]!}`]) {
          const bound = kv.hget(bindingsKey, field);
          if (bound !== null) out.bindings[field] = bound;
          const entry = kv.hget(epochsKey, field);
          if (entry !== null) out.epochs[field] = entry;
        }
      }
      return JSON.stringify(out);
    }
    const entries = JSON.parse(argv[2]!) as {
      bindings: Record<string, string>;
      epochs: Record<string, string>;
    };
    for (const [field, value] of Object.entries(entries.bindings)) {
      let cur = kv.hget(bindingsKey, field);
      if (cur === null) {
        kv.hset(bindingsKey, field, value);
        cur = value;
      }
      const entry = entries.epochs[field];
      if (cur === value && entry !== undefined) {
        if (kv.hget(epochsKey, field) === null) {
          kv.hset(epochsKey, field, entry);
        }
      }
    }
    return 1;
  },
};

/**
 * KEYS: workers, ns_workers, bindings, bindepochs
 * ARGV: now(ms), leaseTtlMs
 * Bucket sweep: prunes long-dead workers (and their unreferenced bindings and
 * epoch entries), rewrites every ns_workers field against the survivors, and
 * returns the clientIds seen, pruned ones included, for drain-config GC.
 */
export const gcScript: Script.Def = {
  lua: `${HELPERS_LUA}
local now, lease = tonumber(ARGV[1]), tonumber(ARGV[2])
local clientIds = {}
local all = redis.call('HGETALL', KEYS[1])
for i = 1, #all, 2 do
  local rec = parseTable(all[i + 1])
  if rec and type(rec.clientId) == 'string' then
    clientIds[#clientIds + 1] = rec.clientId
  end
end
pruneWorkers(KEYS[1], KEYS[3], KEYS[4], now, lease)
for _, field in ipairs(redis.call('HKEYS', KEYS[2])) do
  rewriteNs(KEYS[1], KEYS[2], field, nil, nil, now)
end
return clientIds
`,
  js: (kv, keys, argv) => {
    const [workersKey, nsWorkersKey, bindingsKey, epochsKey] = [
      keys[0]!,
      keys[1]!,
      keys[2]!,
      keys[3]!,
    ];
    const [now, lease] = [Number(argv[0]!), Number(argv[1]!)];
    const clientIds: string[] = [];
    for (const raw of Object.values(kv.hgetall(workersKey))) {
      const clientId = parseObj(raw)?.["clientId"];
      if (typeof clientId === "string") clientIds.push(clientId);
    }
    pruneWorkers(kv, workersKey, bindingsKey, epochsKey, now, lease);
    for (const field of Object.keys(kv.hgetall(nsWorkersKey))) {
      rewriteNs(kv, workersKey, nsWorkersKey, field, null, null, now);
    }
    return clientIds;
  },
};

/**
 * KEYS: ns_workers, workers
 * ARGV: field (`${hostname}${ns}`), sessionId, op('add'/'remove'), now(ms)
 * Rewrites one ns_workers field, pruning sessionIds whose worker is absent
 * or expired.
 */
export const nsWorkersScript: Script.Def = {
  lua: `${HELPERS_LUA}
local field, sessionId, op = ARGV[1], ARGV[2], ARGV[3]
local now = tonumber(ARGV[4])
local ensure = nil
local remove = nil
if op == 'add' then ensure = sessionId else remove = sessionId end
return rewriteNs(KEYS[2], KEYS[1], field, ensure, remove, now)
`,
  js: (kv, keys, argv) => {
    const [nsWorkersKey, workersKey] = [keys[0]!, keys[1]!];
    const [field, sessionId, op] = [argv[0]!, argv[1]!, argv[2]!];
    const now = Number(argv[3]!);
    return rewriteNs(
      kv,
      workersKey,
      nsWorkersKey,
      field,
      op === "add" ? sessionId : null,
      op === "add" ? null : sessionId,
      now,
    );
  },
};