// Inspired from https://github.com/restatedev/examples/blob/main/typescript/patterns-use-cases/src/priorityqueue/queue.ts

import * as restate from "@restatedev/restate-sdk";
import { Context, object, ObjectContext } from "@restatedev/restate-sdk";

interface QueueItem {
  awakeable: string;
  idempotencyKey?: string;
  priority: number;
  leaseDurationMs: number;
}

interface LegacyQueueState {
  itemsv2: Record<string, GroupState>;
  inFlight: number;
  paused: boolean;
  leases: Record<string, number>;
}

interface QueueState {
  groups: Record<string, GroupState>;
  paused: boolean;
  leases: Record<string, number>;
}

interface GroupState {
  id: string;
  items: QueueItem[];
  lastServedTimestamp: number;
}

export const semaphore = object({
  name: "Semaphore",
  handlers: {
    acquire: restate.handlers.object.exclusive(
      {
        ingressPrivate: true,
      },
      async (
        ctx: ObjectContext<LegacyQueueState>,
        req: {
          awakeableId: string;
          priority: number;
          capacity: number;
          leaseDurationMs: number;
          groupId?: string;
          idempotencyKey?: string;
        },
      ): Promise<boolean> => {
        const state = await getState(ctx);

        if (
          req.idempotencyKey &&
          idempotencyKeyAlreadyExists(state.groups, req.idempotencyKey)
        ) {
          return false;
        }

        req.groupId = req.groupId ?? "__ungrouped__";

        if (state.groups[req.groupId] === undefined) {
          state.groups[req.groupId] = {
            id: req.groupId,
            items: [],
            lastServedTimestamp: await ctx.date.now(),
          };
        }

        state.groups[req.groupId].items.push({
          awakeable: req.awakeableId,
          priority: req.priority,
          idempotencyKey: req.idempotencyKey,
          leaseDurationMs: req.leaseDurationMs,
        });

        await tick(ctx, state, req.capacity);

        setState(ctx, state);
        return true;
      },
    ),

    release: restate.handlers.object.exclusive(
      {
        ingressPrivate: true,
      },
      async (
        ctx: ObjectContext<LegacyQueueState>,
        req: {
          leaseId: string;
          capacity: number;
        },
      ): Promise<void> => {
        const state = await getState(ctx);
        delete state.leases[req.leaseId];
        await tick(ctx, state, req.capacity);
        setState(ctx, state);
      },
    ),
    pause: restate.handlers.object.exclusive(
      {},
      async (ctx: ObjectContext<LegacyQueueState>): Promise<void> => {
        const state = await getState(ctx);
        state.paused = true;
        setState(ctx, state);
      },
    ),
    resume: restate.handlers.object.exclusive(
      {},
      async (ctx: ObjectContext<LegacyQueueState>): Promise<void> => {
        const state = await getState(ctx);
        state.paused = false;
        await tick(ctx, state, 1);
        setState(ctx, state);
      },
    ),
    resetInflight: restate.handlers.object.exclusive(
      {},
      async (ctx: ObjectContext<LegacyQueueState>): Promise<void> => {
        const state = await getState(ctx);
        state.leases = {};
        setState(ctx, state);
      },
    ),
    tick: restate.handlers.object.exclusive(
      {},
      async (ctx: ObjectContext<LegacyQueueState>): Promise<void> => {
        const state = await getState(ctx);
        await tick(ctx, state, 1);
        setState(ctx, state);
      },
    ),
  },
  options: {
    journalRetention: 0,
  },
});

// Lower numbers represent higher priority, mirroring Liteque’s semantics.
function selectAndPopItem(
  state: QueueState,
  now: number,
): {
  item: QueueItem;
  groupId: string;
} {
  let selected: {
    priority: number;
    groupId: string;
    index: number;
    groupLastServedTimestamp: number;
  } = {
    priority: Number.MAX_SAFE_INTEGER,
    groupId: "",
    index: 0,
    groupLastServedTimestamp: 0,
  };

  for (const [groupId, group] of Object.entries(state.groups)) {
    for (const [i, item] of group.items.entries()) {
      if (item.priority < selected.priority) {
        selected.priority = item.priority;
        selected.groupId = groupId;
        selected.index = i;
        selected.groupLastServedTimestamp = group.lastServedTimestamp;
      } else if (item.priority === selected.priority) {
        if (group.lastServedTimestamp < selected.groupLastServedTimestamp) {
          selected.priority = item.priority;
          selected.groupId = groupId;
          selected.index = i;
          selected.groupLastServedTimestamp = group.lastServedTimestamp;
        }
      }
    }
  }

  const [item] = state.groups[selected.groupId].items.splice(selected.index, 1);
  state.groups[selected.groupId].lastServedTimestamp = now;
  if (state.groups[selected.groupId].items.length === 0) {
    delete state.groups[selected.groupId];
  }
  return { item, groupId: selected.groupId };
}

function pruneExpiredLeases(state: QueueState, now: number) {
  for (const [leaseId, expiry] of Object.entries(state.leases)) {
    if (expiry <= now) {
      delete state.leases[leaseId];
    }
  }
  return Object.keys(state.leases).length;
}

async function tick(
  ctx: ObjectContext<LegacyQueueState>,
  state: QueueState,
  capacity: number,
): Promise<void> {
  let activeLeases = pruneExpiredLeases(state, await ctx.date.now());
  while (
    !state.paused &&
    activeLeases < capacity &&
    Object.keys(state.groups).length > 0
  ) {
    const now = await ctx.date.now();
    const { item } = selectAndPopItem(state, now);
    state.leases[item.awakeable] = now + item.leaseDurationMs;
    activeLeases++;
    ctx.resolveAwakeable(item.awakeable);
  }
}

async function getState(
  ctx: ObjectContext<LegacyQueueState>,
): Promise<QueueState> {
  const groups = (await ctx.get("itemsv2")) ?? {};
  const paused = (await ctx.get("paused")) ?? false;
  const leases = (await ctx.get("leases")) ?? {};

  return {
    groups,
    paused,
    leases,
  };
}

function idempotencyKeyAlreadyExists(
  items: Record<string, GroupState>,
  key: string,
) {
  for (const group of Object.values(items)) {
    if (group.items.some((item) => item.idempotencyKey === key)) {
      return true;
    }
  }
  return false;
}

function setState(ctx: ObjectContext<LegacyQueueState>, state: QueueState) {
  ctx.set("itemsv2", state.groups);
  ctx.set("leases", state.leases);
  ctx.set("paused", state.paused);
}

export class RestateSemaphore {
  constructor(
    private readonly ctx: Context,
    private readonly id: string,
    private readonly capacity: number,
    private readonly leaseDurationMs: number,
  ) {}

  async acquire(priority: number, groupId?: string, idempotencyKey?: string) {
    const awk = this.ctx.awakeable();
    const res = await this.ctx
      .objectClient<typeof semaphore>({ name: "Semaphore" }, this.id)
      .acquire({
        awakeableId: awk.id,
        priority,
        capacity: this.capacity,
        leaseDurationMs: this.leaseDurationMs,
        groupId,
        idempotencyKey,
      });

    if (!res) {
      return false;
    }

    try {
      await awk.promise;
    } catch (e) {
      if (e instanceof restate.CancelledError) {
        await this.release(awk.id);
      }
      throw e;
    }
    return awk.id;
  }
  async release(leaseId: string) {
    await this.ctx
      .objectClient<typeof semaphore>({ name: "Semaphore" }, this.id)
      .release({ leaseId, capacity: this.capacity });
  }
}