import WebSocket from "ws";
import { JetstreamEventSchema } from "./types.js";
import type { JetstreamConfig, JetstreamEvent } from "./types.js";

const DEFAULT_ENDPOINT = "wss://jetstream2.us-east.bsky.network/subscribe";
const DEFAULT_RECONNECT_DELAY_MS = 5000;
const DEFAULT_CURSOR_BUFFER_US = 5_000_000; // 5 seconds in microseconds (Jetstream cursors use microsecond Unix timestamps)

// --- URL building ---

export const buildJetstreamUrl = (config: JetstreamConfig = {}): string => {
  const url = new URL(config.endpoint ?? DEFAULT_ENDPOINT);

  if (config.wantedCollections?.length) {
    for (const c of config.wantedCollections) {
      url.searchParams.append("wantedCollections", c);
    }
  }

  if (config.wantedDids?.length) {
    for (const d of config.wantedDids) {
      url.searchParams.append("wantedDids", d);
    }
  }

  if (config.requireHello) {
    url.searchParams.set("requireHello", "true");
  }

  if (config.cursor != null) {
    url.searchParams.set("cursor", String(config.cursor));
  }

  if (config.compress) {
    url.searchParams.set("compress", "true");
  }

  if (config.maxMessageSizeBytes != null && config.maxMessageSizeBytes > 0) {
    url.searchParams.set(
      "maxMessageSizeBytes",
      String(config.maxMessageSizeBytes)
    );
  }

  return url.toString();
};

// --- Connection ---

export const createJetstreamConnection = (
  config: JetstreamConfig = {}
): WebSocket => new WebSocket(buildJetstreamUrl(config));

// --- Options update ---

export const buildOptionsUpdate = (
  collections?: readonly string[],
  dids?: readonly string[],
  maxMessageSizeBytes?: number
): string =>
  JSON.stringify({
    type: "options_update",
    payload: {
      ...(collections && { wantedCollections: collections }),
      ...(dids && { wantedDids: dids }),
      ...(maxMessageSizeBytes != null &&
        maxMessageSizeBytes > 0 && { maxMessageSizeBytes }),
    },
  });

export const sendOptionsUpdate = (
  ws: WebSocket,
  collections?: readonly string[],
  dids?: readonly string[],
  maxMessageSizeBytes?: number
): boolean => {
  if (ws.readyState !== WebSocket.OPEN) return false;
  ws.send(buildOptionsUpdate(collections, dids, maxMessageSizeBytes));
  return true;
};

// --- Utilities ---

export const waitForOpen = (ws: WebSocket): Promise<void> =>
  new Promise((resolve) => ws.on("open", resolve));

export const onMessage = (
  ws: WebSocket,
  handler: (event: JetstreamEvent) => void
): void => {
  ws.on("message", (data: Buffer) => {
    try {
      const json = JSON.parse(data.toString());
      const result = JetstreamEventSchema.safeParse(json);
      if (result.success) handler(result.data);
    } catch {}
  });
};

export const attachLifecycle = (
  ws: WebSocket,
  options: {
    onOpen?: () => void;
    onReconnect: () => void;
    reconnectDelay?: number;
  }
): void => {
  if (options.onOpen) ws.on("open", options.onOpen);
  ws.on("error", (err) => console.error("WS error:", err.message));
  ws.on("close", () => {
    const delayMs = options.reconnectDelay ?? DEFAULT_RECONNECT_DELAY_MS;
    console.log(`Disconnected. Reconnecting in ${delayMs / 1000}s...`);
    setTimeout(options.onReconnect, delayMs);
  });
};

// --- Cursor tracking ---

export type CursorRef = { current: number | undefined };

export const createCursorRef = (): CursorRef => ({ current: undefined });

export const onMessageWithCursor = (
  ws: WebSocket,
  handler: (event: JetstreamEvent) => void,
  cursorRef: CursorRef
): void => {
  ws.on("message", (data: Buffer) => {
    try {
      const json = JSON.parse(data.toString());
      const result = JetstreamEventSchema.safeParse(json);
      if (result.success) {
        cursorRef.current = result.data.time_us;
        handler(result.data);
      }
    } catch {}
  });
};

export const attachLifecycleWithCursor = (
  ws: WebSocket,
  options: {
    onOpen?: () => void;
    onReconnect: (cursor: number | undefined) => void;
    reconnectDelay?: number;
    cursorRef: CursorRef;
    cursorBufferUs?: number;
  }
): void => {
  if (options.onOpen) ws.on("open", options.onOpen);
  ws.on("error", (err) => console.error("WS error:", err.message));
  ws.on("close", () => {
    const delayMs = options.reconnectDelay ?? DEFAULT_RECONNECT_DELAY_MS;
    console.log(`Disconnected. Reconnecting in ${delayMs / 1000}s...`);
    const bufferUs = options.cursorBufferUs ?? DEFAULT_CURSOR_BUFFER_US;
    const cursor = options.cursorRef.current
      ? options.cursorRef.current - bufferUs
      : undefined;
    setTimeout(() => options.onReconnect(cursor), delayMs);
  });
};

// --- Stream with reconnect ---

export type StreamOptions = {
  config?: Omit<JetstreamConfig, "cursor">;
  onEvent: (event: JetstreamEvent) => void;
  onOpen?: () => void;
  reconnectDelay?: number;
  cursorBufferUs?: number;
};

export const startStreamWithReconnect = (options: StreamOptions): {
  cursorRef: CursorRef;
  getWs: () => WebSocket;
} => {
  const cursorRef = createCursorRef();
  let currentWs: WebSocket;

  const connect = (cursor?: number) => {
    const ws = createJetstreamConnection({ ...options.config, cursor });
    currentWs = ws;
    onMessageWithCursor(ws, options.onEvent, cursorRef);
    attachLifecycleWithCursor(ws, {
      onOpen: options.onOpen,
      onReconnect: connect,
      reconnectDelay: options.reconnectDelay,
      cursorRef,
      cursorBufferUs: options.cursorBufferUs,
    });
  };

  connect();
  return { cursorRef, getWs: () => currentWs };
};