import { BlockingAtom } from "./blocking-atom.ts";
import { BlockingQueue } from "./blocking-queue.ts";
import { CancellationError } from "./errors.ts";
import { makeGate } from "./gate.ts";

const STREAM_CONFIG_DEFAULT = {
  signal: undefined as AbortSignal | undefined,
  queueSize: 1000,
  threshPause: 800,
  threshResume: 200,
  onPause: () => {},
  onResume: () => {},
};

type StreamEvent =
  | ["yield", MessageEvent["data"]]
  | ["end"]
  | ["error", Error];

export async function* streamSocket(
  ws: WebSocket,
  config_: Partial<typeof STREAM_CONFIG_DEFAULT> = {},
): AsyncGenerator<unknown> {
  const { signal, ...config } = { ...STREAM_CONFIG_DEFAULT, ...config_ };
  signal?.throwIfAborted();

  const queue = new BlockingQueue<StreamEvent>(config.queueSize);
  const gate = makeGate();

  const onMessage = gate.many((m: MessageEvent) => {
    queue.enqueue(["yield", m.data]);

    if (queue.size > config.threshPause && config.onPause) {
      queueMicrotask(config.onPause);
    }
  });

  const onError = gate.once((e: Error) => {
    if (e.message === "Unexpected EOF") {
      // todo: why are we getting this on client shutdown instead of onClose?
      queue.enqueue(["end"]);
    } else {
      queue.enqueue(["error", e]);
    }
  });

  const onClose = gate.once(() => queue.enqueue(["end"]));

  try {
    ws.addEventListener("message", onMessage);
    ws.addEventListener("error", onError);
    ws.addEventListener("close", onClose);

    while (true) {
      signal?.throwIfAborted();

      const [event, value] = await queue.dequeue(signal);
      if (queue.size < config.threshResume && config.onResume) {
        queueMicrotask(config.onResume);
      }

      switch (event) {
        case "yield":
          yield value;
          continue;
        case "error":
          throw value;
        case "end":
          return;
      }
    }
  } finally {
    ws.removeEventListener("message", onMessage);
    ws.removeEventListener("error", onError);
    ws.removeEventListener("close", onClose);
  }
}

export async function takeSocket(
  ws: WebSocket,
  signal?: AbortSignal,
): Promise<MessageEvent["data"]> {
  signal?.throwIfAborted();

  const atom = new BlockingAtom<MessageEvent["data"]>();
  const gate = makeGate();

  const onMessage = gate.once((m: MessageEvent) => atom.set(m.data));
  try {
    ws.addEventListener("message", onMessage);
    const data = await atom.get(signal);
    if (!data) {
      throw new CancellationError("socket read aborted");
    }

    return data;
  } finally {
    ws.removeEventListener("message", onMessage);
  }
}