import {
  ClientMessage,
  parseServerMessage,
  ServerMessage,
} from "./protocol.js";

const CLOSE_NORMAL = 1000;
const CLOSE_NO_STATUS = 1005;

type PromisePair<T> = { promise: Promise<T>; resolve: (value: T) => void };

/**
 * The various states our WebSocket can be in:
 *
 * - "disconnected": We don't have a WebSocket, but plan to create one.
 * - "connecting": We have created the WebSocket and are waiting for the
 *   `onOpen` callback.
 * - "ready": We have an open WebSocket.
 * - "closing": We called `.close()` on the WebSocket and are waiting for the
 *   `onClose` callback before we schedule a reconnect.
 * - "stopping": The application decided to totally stop the WebSocket. We are
 *    waiting for the `onClose` callback before we consider this WebSocket stopped.
 * - "stopped": We have stopped the WebSocket and will never create a new one.
 * 
 * 
                     WebSocket State Machine

                         ┌─────────────┐         ┌─────────────┐
       ┌────onclose─────▶│disconnected │─stop()─▶│   stopped   │
       │                 └─────────────┘         └─────────────┘
       │                        │                       ▲
       │                 new WebSocket()                │
       │                        │                   onclose
       │                        ▼                       │
┌─────────────┐          ┌─────────────┐         ┌─────────────┐
│   closing   │◀─close()─│ connecting  │─stop()─▶│  stopping   │
└─────────────┘          └─────────────┘         └─────────────┘
   │   ▲                        │                       ▲  ▲
   │   │                     onopen                     │  │
   │   │                        │                       │  │
   │   │                        ▼                       │  │
   │   │                 ┌─────────────┐                │  │
   │   └─────────close()─│    ready    │─stop()─────────┘  │
   │                     └─────────────┘                   │
   └─────────────────────────────────────stop()────────────┘
 */
type Socket =
  | { state: "disconnected" }
  | { state: "connecting"; ws: WebSocket }
  | { state: "ready"; ws: WebSocket }
  | { state: "closing" }
  | { state: "stopping"; promisePair: PromisePair<null> }
  | { state: "stopped" };

function promisePair<T>(): PromisePair<T> {
  let resolvePromise: (value: T) => void;
  const promise = new Promise<T>(resolve => {
    resolvePromise = resolve;
  });
  return { promise, resolve: resolvePromise! };
}

/**
 * A wrapper around a websocket that handles errors, reconnection, and message
 * parsing.
 */
export class WebSocketManager {
  private socket: Socket;

  /** Upon HTTPS/WSS failure, the first jittered backoff duration, in ms. */
  private readonly initialBackoff: number;

  /** We backoff exponentially, but we need to cap that--this is the jittered max. */
  private readonly maxBackoff: number;

  /** How many times have we failed consecutively? */
  private retries: number;

  private readonly uri: string;
  private readonly onOpen: () => void;
  private readonly onMessage: (message: ServerMessage) => void;
  private readonly webSocketConstructor: typeof WebSocket;

  constructor(
    uri: string,
    onOpen: () => void,
    onMessage: (message: ServerMessage) => void,
    webSocketConstructor: typeof WebSocket
  ) {
    this.webSocketConstructor = webSocketConstructor;
    this.socket = { state: "disconnected" };
    this.initialBackoff = 100;
    this.maxBackoff = 16000;
    this.retries = 0;

    this.uri = uri;
    this.onOpen = onOpen;
    this.onMessage = onMessage;

    this.connect();
  }

  private async connect() {
    if (
      this.socket.state === "closing" ||
      this.socket.state === "stopping" ||
      this.socket.state === "stopped"
    ) {
      return;
    }
    if (this.socket.state !== "disconnected") {
      throw new Error("Didn't start connection from disconnected state");
    }

    const ws = new this.webSocketConstructor(this.uri);
    this.socket = {
      state: "connecting",
      ws,
    };
    ws.onopen = () => {
      if (this.socket.state !== "connecting") {
        throw new Error("onopen called with socket not in connecting state");
      }
      this.socket = { state: "ready", ws };
      this.onOpen();
    };
    // NB: The WebSocket API calls `onclose` even if connection fails, so we can route all error paths through `onclose`.
    ws.onerror = error => {
      const message = (<ErrorEvent>error).message;
      console.log(`WebSocket error: ${message}`);
      this.closeAndReconnect();
    };
    ws.onmessage = message => {
      // TODO(CX-1498): We reset the retry counter on any successful message.
      // This is not ideal and we should improve this further.
      this.retries = 0;
      const serverMessage = parseServerMessage(JSON.parse(message.data));
      this.onMessage(serverMessage);
    };
    ws.onclose = event => {
      if (event.code !== CLOSE_NORMAL && event.code !== CLOSE_NO_STATUS) {
        let msg = `WebSocket closed unexpectedly with code ${event.code}`;
        if (event.reason) {
          msg += `: ${event.reason}`;
        }
        console.error(msg);
      }
      if (this.socket.state === "stopping") {
        this.socket.promisePair.resolve(null);
        this.socket = { state: "stopped" };
        return;
      }
      this.socket = { state: "disconnected" };
      const backoff = this.nextBackoff();
      console.log(`Attempting reconnect in ${backoff}ms`);
      setTimeout(() => this.connect(), backoff);
    };
  }

  sendMessage(message: ClientMessage) {
    if (this.socket.state === "ready") {
      const request = JSON.stringify(message);
      try {
        this.socket.ws.send(request);
      } catch (error: any) {
        console.log(
          `Failed to send message on WebSocket, reconnecting: ${error}`
        );
        this.closeAndReconnect();
      }
    }
  }

  /**
   * Close the WebSocket and schedule a reconnect when it completes closing.
   *
   * This should be used when we hit an error and would like to restart the session.
   */
  private closeAndReconnect() {
    switch (this.socket.state) {
      case "disconnected":
      case "closing":
      case "stopping":
      case "stopped":
        // Nothing to do if we don't have a WebSocket.
        return;
      case "connecting":
      case "ready":
        this.socket.ws.close();
        this.socket = {
          state: "closing",
        };
        return;
      default: {
        // Enforce that the switch-case is exhaustive.
        // eslint-disable-next-line  @typescript-eslint/no-unused-vars
        const _: never = this.socket;
      }
    }
  }

  /**
   * Close the WebSocket and never reconnect.
   * @returns A Promise that resolves when the WebSocket `onClose` callback is called.
   */
  async stop(): Promise<void> {
    switch (this.socket.state) {
      case "stopped":
        return;
      case "connecting":
      case "ready":
        this.socket.ws.close();
        this.socket = {
          state: "stopping",
          promisePair: promisePair(),
        };
        await this.socket.promisePair.promise;
        return;
      case "closing":
        // We're already closing the WebSocket, so just upgrade the state
        // to "stopping" so we don't reconnect.
        this.socket = {
          state: "stopping",
          promisePair: promisePair(),
        };
        await this.socket.promisePair.promise;
        return;
      case "disconnected":
        // We're disconnected so switch the state to "stopped" so the reconnect
        // timeout doesn't create a new WebSocket.
        this.socket = { state: "stopped" };
        return;
      case "stopping":
        await this.socket.promisePair.promise;
        return;
      default: {
        // Enforce that the switch-case is exhaustive.
        const _: never = this.socket;
      }
    }
  }

  private nextBackoff(): number {
    const baseBackoff = this.initialBackoff * Math.pow(2, this.retries);
    this.retries += 1;
    const actualBackoff = Math.min(baseBackoff, this.maxBackoff);
    const jitter = actualBackoff * (Math.random() - 0.5);
    return actualBackoff + jitter;
  }
}