type StreamPayload = unknown;
type StreamSubscriber = (
  text: string,
  finished: boolean,
  payload: StreamPayload,
) => void;

interface StreamState {
  text: string;
  finished: boolean;
  payload: StreamPayload;
  subs: Set<StreamSubscriber>;
}

const streams: Record<string, StreamState> = {};

function createEmptyStream(): StreamState {
  return { text: "", finished: false, payload: null, subs: new Set() };
}

function notify(stream: StreamState): void {
  stream.subs.forEach((cb) => cb(stream.text, stream.finished, stream.payload));
}

function ensureStream(id: string): StreamState {
  streams[id] = streams[id] || createEmptyStream();
  return streams[id];
}

function openStream(
  id: string,
  initial = "",
  finished = false,
  payload: StreamPayload = null,
): void {
  if (!id) return;
  const stream = ensureStream(id);
  stream.text = initial || "";
  stream.finished = Boolean(finished);
  stream.payload = payload || null;
  notify(stream);
}

function updateStream(
  id: string,
  text: string,
  finished = false,
  payload: StreamPayload = null,
): void {
  if (!id || !streams[id]) return;
  streams[id].text = text || "";
  streams[id].finished = Boolean(finished);
  streams[id].payload = payload || null;
  notify(streams[id]);
}

function subscribe(id: string, cb: StreamSubscriber): () => void {
  if (!id) return () => {};
  const stream = ensureStream(id);
  stream.subs.add(cb);
  cb(stream.text, stream.finished, stream.payload);

  return () => {
    const current = streams[id];
    if (!current) return;
    current.subs.delete(cb);
    if (current.subs.size === 0 && current.finished) {
      delete streams[id];
    }
  };
}

function closeStream(id: string): void {
  if (!id || !streams[id]) return;
  streams[id].subs.forEach((cb) => cb(streams[id].text, true, streams[id].payload));
  delete streams[id];
}

function getSnapshot(id: string): Omit<StreamState, "subs"> {
  if (!id || !streams[id]) return { text: "", finished: false, payload: null };
  return {
    text: streams[id].text,
    finished: streams[id].finished,
    payload: streams[id].payload,
  };
}

export default {
  openStream,
  updateStream,
  subscribe,
  closeStream,
  getSnapshot,
};
