// Derivation helpers over NodeInventory snapshots. These are pure
// projections — no state of their own — so the placement algorithm can
// remain a pure function over what the worker advertised.

import { Capability, CapabilityRequirement, GpuResource, JobRequirements, NodeInventory } from "./types";

/**
 * Free capacity available for placement. `(totalCapacity * utilizationCap)
 * - reservedCapacity`. Can be negative if the node is over-allocated;
 * placement should treat negative values as "no room".
 *
 * @public
 */
export function freeCapacity(inv: NodeInventory): number {
  return inv.totalCapacity * inv.utilizationCap - inv.reservedCapacity;
}

/** @public */
export function freeCores(inv: NodeInventory): number {
  return inv.totalCores - inv.reservedCores;
}

/**
 * Fraction (0-1) of the node's total capacity currently reserved.
 * Used by the bin-pack scoring strategy.
 *
 * @public
 */
export function usedCapacityFraction(inv: NodeInventory): number {
  if (inv.totalCapacity <= 0) return 0;
  return inv.reservedCapacity / inv.totalCapacity;
}

/** Two attribute sets are provably disjoint when some shared key carries
 *  different values — "cardIndex 0" cannot be "cardIndex 1". Sets sharing
 *  no keys are NOT disjoint: they might describe the same device. */
function attributesDisjoint(
  a: Record<string, string> | undefined,
  b: Record<string, string> | undefined
): boolean {
  if (!a || !b) return false;
  for (const [k, v] of Object.entries(a)) {
    const other = b[k];
    if (other !== undefined && other !== v) return true;
  }
  return false;
}

/** Does an advertised capability satisfy a requirement's attribute filter?
 *  Every requested key must be present with the requested value. */
function entryMatches(
  cap: Capability,
  attributeMatches: Record<string, string> | undefined
): boolean {
  if (!attributeMatches) return true;
  for (const [k, v] of Object.entries(attributeMatches)) {
    if (cap.attributes?.[k] !== v) return false;
  }
  return true;
}

/**
 * Free count of the capability subset a requirement names: total across
 * every advertised entry matching the requirement's attributes, minus
 * every reservation not provably disjoint from them. The overlap rule is
 * deliberately conservative — a reservation whose attributes share no key
 * with the query could have consumed any entry, so it subtracts. Placement
 * may under-place on heterogeneous partial attribute sets; it never
 * double-books a device.
 *
 * @public
 */
export function freeCapabilityFor(
  inv: NodeInventory,
  req: CapabilityRequirement
): number {
  let total = 0;
  for (const c of inv.capabilities) {
    if (c.name === req.name && entryMatches(c, req.attributeMatches)) total += c.count;
  }
  let used = 0;
  for (const r of inv.reservedCapabilities) {
    if (r.name === req.name && !attributesDisjoint(r.attributes, req.attributeMatches))
      used += r.count;
  }
  return Math.max(0, total - used);
}

/** The host sockets already claimed on a node, as a set of canonical
 *  `PORT/proto` tokens. @public */
export function reservedHostPortSet(inv: NodeInventory): Set<string> {
  return new Set(inv.reservedHostPorts ?? []);
}

/** The capability a worker advertises when its job socket is listening;
 *  the handover strategy cannot run without it. A FACILITY, not a resource:
 *  every instance on the worker shares the one listener, so a job requiring
 *  it must match it but never consume it — reserving the worker's single
 *  count would let exactly one handover-capable job onto each worker.
 *
 *  @public */
export const JOB_SOCKET_CAPABILITY = "norsk.io/job-socket";

/** Capabilities that are matched but never reserved (see JOB_SOCKET_CAPABILITY). */
const SHARED_CAPABILITIES: ReadonlySet<string> = new Set([JOB_SOCKET_CAPABILITY]);

/**
 * One placed job's claim on a node: its typed requirements plus the GPU
 * device placement chose for it (`PlacementResult.gpuIndex`), when it
 * needs one and the device is known. A GPU claim with no index cannot be
 * attributed to a device and reserves none.
 *
 * @public
 */
export interface PlacedClaim {
  requirements: JobRequirements;
  gpuIndex?: number;
}

/** @public */
export interface GpuReservation {
  index: number;
  reservedCapacity: number;
}

/** What `deriveReservation` sums. @public */
export interface DerivedReservation {
  capacity: number;
  cores: number;
  capabilities: Capability[];
  hostPorts: string[];
  /** Per device index, only devices with a claim on them. */
  gpus: GpuReservation[];
}

/** A requirement number outside its domain (negative, NaN, infinite)
 *  contributes nothing: summed as-is, a negative capacity would FREE the
 *  node. validateBundle rejects these at submit; this is the backstop for
 *  a bundle that reached us another way. */
function nonNegative(v: number | undefined): number {
  return v !== undefined && Number.isFinite(v) && v > 0 ? v : 0;
}

/**
 * Sum a set of placed jobs' claims into a node reservation — the single
 * derivation AutoManager overlays onto wire inventory (which arrives with
 * zeroed reservation fields today). Capability reservations preserve the
 * requirement's attribute filter, so an identity-bound claim ("decklink
 * cardIndex 0") blocks that card and not its siblings; host ports union
 * into the claimed-socket set; GPU capacity is attributed to the device
 * each job was placed on.
 *
 * @public
 */
export function deriveReservation(claims: PlacedClaim[]): DerivedReservation {
  let capacity = 0;
  let cores = 0;
  const caps = new Map<string, { name: string; attributes?: Record<string, string>; count: number }>();
  const ports = new Set<string>();
  const gpus = new Map<number, number>();
  for (const { requirements: req, gpuIndex } of claims) {
    capacity += nonNegative(req.requiredCapacity);
    cores += nonNegative(req.requiredCores);
    for (const c of req.requiredCapabilities) {
      if (SHARED_CAPABILITIES.has(c.name)) continue;
      const attrKey = c.attributeMatches
        ? JSON.stringify(Object.entries(c.attributeMatches).sort(([a], [b]) => a.localeCompare(b)))
        : "";
      const key = `${c.name} ${attrKey}`;
      const existing = caps.get(key);
      if (existing) existing.count += c.count;
      else
        caps.set(key, {
          name: c.name,
          ...(c.attributeMatches ? { attributes: c.attributeMatches } : {}),
          count: c.count,
        });
    }
    for (const p of req.requiredHostPorts ?? []) ports.add(p);
    const gpuCapacity = nonNegative(req.requiredGpuCapacity);
    if (gpuCapacity > 0 && gpuIndex !== undefined) {
      gpus.set(gpuIndex, (gpus.get(gpuIndex) ?? 0) + gpuCapacity);
    }
  }
  const capabilities: Capability[] = [...caps.values()].map((c) => ({
    name: c.name,
    count: c.count,
    ...(c.attributes ? { attributes: c.attributes } : {}),
  }));
  return {
    capacity,
    cores,
    capabilities,
    hostPorts: [...ports],
    gpus: [...gpus.entries()].map(([index, reservedCapacity]) => ({ index, reservedCapacity })),
  };
}

/**
 * Overlay a derived reservation onto a wire inventory, deferring to a
 * non-zero wire value — if a producer ever advertises real reservations
 * they are authoritative. Per GPU device the same rule: a device the wire
 * reports as reserved (any capacity, or exclusively) is kept; otherwise
 * the derived capacity is written on, and a device the claims fill
 * (reserved >= totalCapacity) is `exclusivelyReserved` — the whole-GPU
 * job of the `GpuResource` doc.
 *
 * @public
 */
export function overlayReservation(inv: NodeInventory, reservation: DerivedReservation): NodeInventory {
  const derivedGpu = new Map(reservation.gpus.map((g) => [g.index, g.reservedCapacity]));
  const gpus: GpuResource[] = inv.gpus.map((g) => {
    if (g.reservedCapacity > 0 || g.exclusivelyReserved) return g;
    const reserved = derivedGpu.get(g.index);
    if (reserved === undefined) return g;
    return {
      ...g,
      reservedCapacity: reserved,
      exclusivelyReserved: reserved >= g.totalCapacity,
    };
  });
  return {
    ...inv,
    gpus,
    reservedCapacity: inv.reservedCapacity > 0 ? inv.reservedCapacity : reservation.capacity,
    reservedCores: inv.reservedCores > 0 ? inv.reservedCores : reservation.cores,
    reservedCapabilities:
      inv.reservedCapabilities.length > 0 ? inv.reservedCapabilities : reservation.capabilities,
    reservedHostPorts:
      inv.reservedHostPorts && inv.reservedHostPorts.length > 0
        ? inv.reservedHostPorts
        : reservation.hostPorts,
  };
}