import type { Command } from "./operations";
import { xformMany } from "./transformation";
import { assert } from "./utils";

type OpsAckMessage = {
  type: "server";
  ops?: Command[];
  ack?: boolean;
  seq: number;
};

type LoadedMessage = { type: "loaded" };

type ErrorMessage = {
  type: "error";
  message: string;
};

export type ServerMessage = OpsAckMessage | LoadedMessage | ErrorMessage;

export type ClientMessage = {
  type: "client";
  ops: Command[];
  fromSeq: number;
};

// Clients connect to workspaces by WebSocket, as implemented by @yakatak/app
// for /api/workspaces/[id]/ws.
// TODO expose whether there are outstanding requests
export class WebSocketOTClient {
  // Ready for pollServerOperations() to be called.
  onPollReady?: () => void;
  // Initial batch of server operations have been received.
  onLoaded?: () => void;

  private ws: WebSocket | null = null;
  private error: string | null = null;
  private url: string;
  private clientId: string;
  private reconnectAttempt = 0;

  // Unpolled server operations.
  private server: Command[] = [];

  // The last known committed sequence number.
  private seq: number;
  // Local operations that have not yet been acked.
  private local: Command[];
  // Number of local operations that have already been sent; we use the
  // approach of Google Wave and only have one outstanding request at a time.
  private numSent;

  constructor(
    url: string,
    clientId: string,
    { seq, local, numSent } = { seq: -1, local: [] as Command[], numSent: 0 },
  ) {
    this.url = url;
    this.clientId = clientId;
    this.seq = seq;
    this.local = local;
    this.numSent = numSent;
  }

  connect() {
    assert(!this.ws);
    const url = `${this.url}?client=${this.clientId}`;
    const ws = new WebSocket(url);
    ws.onopen = () => {
      this.ws = ws;
      this.reconnectAttempt = 0;
      // FIXME server should discard idempotent operation on reconnect
      this.maybeSendOps();
    };
    ws.onmessage = (e) => this.onMessage(e);
    // WebSocket errors are generally unhelpful, in part due to security
    // restrictions. Mostly ignore and rely on the close code.
    ws.onerror = (e) => console.warn("WebSocket error");
    ws.onclose = (e) => this.scheduleReconnect(e);
  }

  private exponentialBackoff() {
    const base = Math.min(500 * 2 ** this.reconnectAttempt++, 30_000);
    const jitter = Math.random() + 0.5;
    return Math.round(base * jitter);
  }

  private scheduleReconnect(e: CloseEvent) {
    this.ws = null;
    // Normal and abnormal closure codes.
    if (e.code === 1000 || e.code === 1006) {
      if (this.error == null) {
        const delay = this.exponentialBackoff();
        console.log(`Reconnecting WebSocket in ${delay}ms`);
        setTimeout(() => this.connect(), delay);
      }
    } else {
      this.setError(`WebSocket close code: ${e.code}`);
    }
  }

  pollServerOperations(): Command[] {
    if (this.error !== null) throw new Error(this.error);
    return this.server.splice(0, this.server.length);
  }

  sendLocalOperations(...commands: Command[]) {
    if (commands.length === 0) return;
    [commands, this.server] = xformMany(commands, this.server);
    this.local.push(...commands);
    this.maybeSendOps();
  }

  private onMessage(event: MessageEvent) {
    const data = JSON.parse(event.data) as ServerMessage;
    if (data.type === "server") {
      if (data.ack) {
        this.local.splice(0, this.numSent);
        this.numSent = 0;
      }
      let server = data.ops ?? [];
      [this.local, server] = xformMany(this.local, server);
      assert(this.seq < data.seq);
      this.seq = data.seq;
      const ready = this.server.length === 0 && server.length > 0;
      this.server.push(...server);
      this.maybeSendOps();
      if (ready) this.onPollReady?.();
    } else if (data.type === "loaded") {
      this.onLoaded?.();
      delete this.onLoaded;
    } else if (data.type === "error") {
      this.setError(data.message);
    } else {
      const _unreachable: never = data;
    }
  }

  private maybeSendOps() {
    if (!this.ws || this.seq < 0 || this.numSent > 0 || this.local.length === 0) return;

    this.numSent = this.local.length;
    this.send({ type: "client", ops: this.local, fromSeq: this.seq });
  }

  private send(msg: ClientMessage) {
    assert(this.ws);
    this.ws.send(JSON.stringify(msg));
  }

  private setError(msg: string) {
    if (this.error == null) {
      this.error = msg;
      this.onPollReady?.();
    }
  }
}