import * as decoding from "lib0/decoding";
import * as encoding from "lib0/encoding";
import * as awarenessProtocol from "y-protocols/awareness";
import * as syncProtocol from "y-protocols/sync";
import * as Y from "yjs";
import {
  CLOSE_DOWNGRADED,
  CLOSE_REAUTH,
  CLOSE_REVOKED,
  EDITOR_WS_PROTOCOL,
  FILES_MAP,
  MESSAGE_AWARENESS,
  MESSAGE_SYNC,
  textKey,
} from "../lib/editor-protocol";
import type { ClientControl, FileMeta, RoomControl } from "../lib/editor-protocol";

export type RoomEvents = {
  onControl(message: RoomControl): void;
  onState(state: "connecting" | "online" | "offline" | "revoked"): void;
  /** Asked for a fresh token after the current one expired or write access changed. */
  refreshToken(): Promise<string>;
};

const MIN_BACKOFF_MS = 500;
const MAX_BACKOFF_MS = 10_000;
const OPEN_TIMEOUT_MS = 15_000;

/**
 * Yjs client for the EditorRoom: one Y.Doc for the whole branch, y-protocols sync and awareness over a
 * single WebSocket, reconnecting with backoff. Edits made while offline merge on reconnect (CRDT).
 */
export class RoomClient {
  readonly doc = new Y.Doc();
  readonly awareness = new awarenessProtocol.Awareness(this.doc);
  readonly files: Y.Map<FileMeta> = this.doc.getMap<FileMeta>(FILES_MAP);
  writable = false;
  head = "";
  private socket: WebSocket | null = null;
  private backoff = MIN_BACKOFF_MS;
  private failedOpens = 0;
  private stopped = false;
  private synced = false;
  private waiters: Array<() => void> = [];

  constructor(
    private readonly url: string,
    private token: string,
    private readonly events: RoomEvents,
  ) {
    this.doc.on("update", (update: Uint8Array, origin: unknown) => {
      if (origin === this) return;
      const encoder = encoding.createEncoder();
      encoding.writeVarUint(encoder, MESSAGE_SYNC);
      syncProtocol.writeUpdate(encoder, update);
      this.send(encoding.toUint8Array(encoder));
    });
    this.awareness.on("update", ({ added, updated, removed }: Record<string, number[]>, origin: unknown) => {
      if (origin === this) return;
      const changed = [...added, ...updated, ...removed];
      const encoder = encoding.createEncoder();
      encoding.writeVarUint(encoder, MESSAGE_AWARENESS);
      encoding.writeVarUint8Array(encoder, awarenessProtocol.encodeAwarenessUpdate(this.awareness, changed));
      this.send(encoding.toUint8Array(encoder));
    });
    this.connect();
  }

  text(path: string): Y.Text {
    return this.doc.getText(textKey(path));
  }

  control(message: ClientControl): void {
    if (this.socket?.readyState === WebSocket.OPEN) this.socket.send(JSON.stringify(message));
  }

  /** Resolves once the room has the file's metadata (seeding it from the branch head if needed). */
  async open(path: string): Promise<FileMeta> {
    await this.whenSynced();
    const existing = this.files.get(path);
    if (existing) return existing;
    return new Promise((resolve, reject) => {
      const timer = setTimeout(() => {
        this.files.unobserve(observer);
        this.missing.delete(path);
        reject(new Error(`Timed out opening ${path}`));
      }, OPEN_TIMEOUT_MS);
      const observer = () => {
        const meta = this.files.get(path);
        if (!meta) return;
        clearTimeout(timer);
        this.files.unobserve(observer);
        this.missing.delete(path);
        resolve(meta);
      };
      this.files.observe(observer);
      this.missing.set(path, (message) => {
        clearTimeout(timer);
        this.files.unobserve(observer);
        this.missing.delete(path);
        reject(new Error(message));
      });
      this.control({ type: "open", path });
    });
  }

  private missing = new Map<string, (message: string) => void>();

  whenSynced(): Promise<void> {
    if (this.synced) return Promise.resolve();
    return new Promise((resolve) => this.waiters.push(resolve));
  }

  dispose(): void {
    this.stopped = true;
    awarenessProtocol.removeAwarenessStates(this.awareness, [this.doc.clientID], "dispose");
    this.socket?.close();
  }

  private send(frame: Uint8Array): void {
    if (this.socket?.readyState === WebSocket.OPEN) this.socket.send(frame);
  }

  private connect(): void {
    if (this.stopped) return;
    this.events.onState("connecting");
    const socket = new WebSocket(this.url, [EDITOR_WS_PROTOCOL, this.token]);
    socket.binaryType = "arraybuffer";
    this.socket = socket;
    let opened = false;
    socket.onopen = () => {
      opened = true;
      this.failedOpens = 0;
      this.backoff = MIN_BACKOFF_MS;
      const encoder = encoding.createEncoder();
      encoding.writeVarUint(encoder, MESSAGE_SYNC);
      syncProtocol.writeSyncStep1(encoder, this.doc);
      socket.send(encoding.toUint8Array(encoder));
      if (this.awareness.getLocalState() !== null) {
        const aw = encoding.createEncoder();
        encoding.writeVarUint(aw, MESSAGE_AWARENESS);
        encoding.writeVarUint8Array(aw, awarenessProtocol.encodeAwarenessUpdate(this.awareness, [this.doc.clientID]));
        socket.send(encoding.toUint8Array(aw));
      }
      this.events.onState("online");
    };
    socket.onmessage = (event) => this.onMessage(event.data);
    socket.onclose = (event) => {
      if (!opened) this.failedOpens++;
      void this.onClose(event.code);
    };
  }

  private onMessage(data: unknown): void {
    if (typeof data === "string") {
      const message = JSON.parse(data) as RoomControl;
      if (message.type === "welcome") {
        this.writable = message.writable;
        this.head = message.head;
      }
      if (message.type === "status") this.head = message.head;
      if (message.type === "error" && message.path) this.missing.get(message.path)?.(message.message);
      this.events.onControl(message);
      return;
    }
    const decoder = decoding.createDecoder(new Uint8Array(data as ArrayBuffer));
    const type = decoding.readVarUint(decoder);
    if (type === MESSAGE_SYNC) {
      const encoder = encoding.createEncoder();
      encoding.writeVarUint(encoder, MESSAGE_SYNC);
      const subtype = syncProtocol.readSyncMessage(decoder, encoder, this.doc, this);
      if (encoding.length(encoder) > 1) this.send(encoding.toUint8Array(encoder));
      if (subtype === syncProtocol.messageYjsSyncStep2 && !this.synced) {
        this.synced = true;
        for (const resolve of this.waiters.splice(0)) resolve();
      }
    } else if (type === MESSAGE_AWARENESS) {
      awarenessProtocol.applyAwarenessUpdate(this.awareness, decoding.readVarUint8Array(decoder), this);
    }
  }

  private async onClose(code: number): Promise<void> {
    this.socket = null;
    const peers = [...this.awareness.getStates().keys()].filter((id) => id !== this.doc.clientID);
    awarenessProtocol.removeAwarenessStates(this.awareness, peers, this);
    if (this.stopped) return;
    if (code === CLOSE_REVOKED) {
      this.stopped = true;
      this.writable = false;
      this.events.onState("revoked");
      return;
    }
    this.events.onState("offline");
    // An upgrade rejected with 401 (e.g. the token expired while offline) surfaces as a plain close.
    if (code === CLOSE_REAUTH || code === CLOSE_DOWNGRADED || this.failedOpens >= 2) {
      try {
        this.token = await this.events.refreshToken();
        this.failedOpens = 0;
      } catch {
        this.events.onState("revoked");
        return;
      }
    }
    setTimeout(() => this.connect(), this.backoff);
    this.backoff = Math.min(this.backoff * 2, MAX_BACKOFF_MS);
  }
}