import { ExponentialBackoff, fromExtendedJson, toExtendedJson } from '@xo-cash/utils';
import type { PrivateKey } from '@xo-cash/primitives';

import { WsMessageSchema, WsSuccessResponseSchema, WsErrorResponseSchema, type WsMessage } from './types.js';
import { BaseSyncClient, type ResourcesRecord } from '../shared/base-client.js';

/** A request sent over the sync server's WebSocket connection. */
type WsRequest = {
  id?: string;
  path: string;
  body?: unknown;
  headers?: Record<string, string>;
};

type PendingRequest = {
  resolve: (body: unknown) => void;
  reject: (error: Error) => void;
};

export type WsClientOptions = {
  /** Called for server-pushed events and uncorrelated server errors. */
  onMessage: (message: WsMessage) => void;

  /** Called for transport errors and malformed server messages. */
  onError: (error: Error) => void;
};

/**
 * Minimal WebSocket client for the XO sync server.
 *
 * - one connection is opened at `/ws`, on demand: every request awaits the
 *   connection first, so callers never need to wait for one themselves;
 * - request IDs correlate concurrent read/write/unsubscribe calls;
 * - subscriptions are remembered and re-sent on every new socket, because the
 *   server drops a socket's topics when it closes; and
 * - a dropped connection fails pending requests and is reopened with backoff
 *   until `disconnect()` is called. Requests made in the meantime wait for it.
 *
 * Heartbeats and request timeouts are not implemented.
 */
export class WsClient<PayloadType = unknown> extends BaseSyncClient<PayloadType> {
  private readonly url: string;
  private readonly privateKey: PrivateKey;
  private readonly pendingRequests = new Map<string, PendingRequest>();
  private readonly messageListeners = new Set<(message: WsMessage) => void>();
  private readonly errorListeners = new Set<(error: Error) => void>();
  private readonly textEncoder = new TextEncoder();
  
  /** Resource IDs to (re)subscribe to whenever a socket opens. */
  private readonly subscriptions = new Set<string>();

  /**
   * The current connection: resolves to an open socket once a connection attempt succeeds.
   * Undefined until `connect()` and after `disconnect()`, which also means "don't reconnect".
   */
  private webSocket: Promise<WebSocket> | undefined;

  /** Retries a failed connection attempt forever (maxAttempts 0), waiting up to 30s between attempts. */
  private readonly connectBackoff = new ExponentialBackoff({ maxAttempts: 0, maxDelay: 30_000 });

  constructor(
    url: string,
    privateKey: PrivateKey,
    options: Partial<WsClientOptions> = {},
  ) {
    super();

    this.url = url;
    this.privateKey = privateKey;
    this.messageListeners.add(options.onMessage ?? (() => {}));
    this.errorListeners.add(options.onError ?? (() => {}));
  }

  /** Connect, and keep reconnecting on drops until `disconnect()` is called. */
  async connect(): Promise<void> {
    await this.getWebSocket();
  }

  /**
   * Close the socket and stop reconnecting.
   *
   * Closing a WebSocket automatically removes all of its server-side topics,
   * so disconnect does not need to send an unsubscribe request first.
   * Subscriptions are kept, and are replayed if the client connects again.
   *
   * Resolves without waiting for the socket to finish closing; `connection-updated` fires when it does.
   */
  async disconnect(): Promise<void> {
    const attempt = this.webSocket;
    this.webSocket = undefined;

    this.rejectPendingRequests(new Error('WebSocket client disconnected'));

    // If the attempt is still retrying, it sees it is no longer current after its next handshake,
    // closes that socket and stops. Its rejection is handled here.
    attempt?.then((socket) => socket.close(), () => {});
  }

  /**
   * The current socket, starting a connection attempt if there isn't one.
   *
   * A new socket has no topics on the server, so starting an attempt also queues a subscribe for
   * every remembered subscription. They are sent through send(), so they go out once the socket opens.
   */
  private getWebSocket(): Promise<WebSocket> {
    if (this.webSocket) {
      return this.webSocket;
    }

    this.webSocket = this.openWebSocket();

    for (const resourceId of this.subscriptions) {
      // A failure here means the attempt was stopped by disconnect() or the socket closed before the
      // frame was sent. A close starts a new attempt, which replays this again, so there is nothing to do.
      this.send('/data/subscribe', { resourceId: [resourceId] }).catch(() => {});
    }

    return this.webSocket;
  }

  /**
   * Retry opening a socket until one opens. Use getWebSocket() instead of calling this directly.
   *
   * `this.webSocket === attempt` tells whether this attempt is still the current connection.
   * It stops being current when `disconnect()` is called.
   */
  private openWebSocket(): Promise<WebSocket> {
    const attempt: Promise<WebSocket> = this.connectBackoff.run(async ({ stopRetries }) => {
      // `attempt` is only read after this await: run() calls this function before `attempt` is assigned.
      const socket = await WsClient.waitForOpen(new WebSocket(WsClient.httpToWsUrl(this.url)));

      // disconnect() was called while connecting. The thrown error rejects `attempt`.
      if (this.webSocket !== attempt) {
        socket.close();
        stopRetries('WebSocket client disconnected');
        throw new Error('WebSocket client disconnected');
      }

      socket.onmessage = (event) => this.handleMessage(String(event.data));
      socket.onerror = () => this.emitError(new Error('WebSocket error'));
      socket.onclose = () => {
        this.rejectPendingRequests(new Error('WebSocket connection closed'));
        this.emit('connection-updated', 'disconnected');

        // Reconnect, unless disconnect() was called
        if (this.webSocket === attempt) {
          this.webSocket = undefined;
          this.getWebSocket();
        }
      };

      this.emit('connection-updated', 'connected');

      return socket;
    }, {
      // Report failed handshakes, but not the error thrown above after disconnect()
      onError: (error) => {
        if (this.webSocket === attempt) {
          this.emitError(error);
        }
      },
    });

    // Nothing may be awaiting this attempt (e.g. a reconnect started by onclose), and it rejects if disconnect() stops it.
    // Mark the rejection as handled so it isn't reported as unhandled. Callers awaiting `attempt` still receive it.
    attempt.catch(() => {});

    return attempt;
  }

  async readMany(resourceIds: string[]): Promise<ResourcesRecord<PayloadType>> {
    // Read the resources from the server
    const response = await this.request('/data/get', { resourceId: resourceIds });

    // Parse the response as a record of resources
    const resources = BaseSyncClient.resourcesResponseSchema.parse(response);

    // Decode the resources and convert the blobs from Uint8Arrays to the type T
    const decodedResources = BaseSyncClient.decodeResourcesResponse<PayloadType>(resources);

    return decodedResources;
  }

  async writeMany(resourcesToWrite: Record<string, PayloadType>): Promise<ResourcesRecord<PayloadType>> {
    // Sign the resources
    const signedResources = Object.entries(resourcesToWrite).map(([resourceId, value]) => {
      const valueStr = toExtendedJson(value);
      const valueBytes = this.textEncoder.encode(valueStr);
      return BaseSyncClient.signWriteRequest(this.privateKey.toBytes(), resourceId, valueBytes);
    });

    // Write the resources to the server
    const response = await this.request('/data/write', { resources: signedResources });

    // Parse the response as a record of resources
    const resources = BaseSyncClient.resourcesResponseSchema.parse(response);

    // Decode the resources and convert the blobs from Uint8Arrays to the type T
    const decodedResources = BaseSyncClient.decodeResourcesResponse<PayloadType>(resources);

    return decodedResources;
  }

  /**
   * Subscribe to one resource topic, now and on every future socket.
   *
   * The server never replies to a successful subscribe (the route stays pending until the topic is
   * removed), so this resolves once the frame is sent. A failed subscribe comes back as an `error`
   * message that matches no pending request, which handleMessage() reports through emitError().
   * 
   * @param resourceId - The resource ID to subscribe to
   */
  async subscribe(resourceId: string): Promise<void> {
    await this.send('/data/subscribe', { resourceId: [resourceId] });

    // Remembered only after sending: if send() started a new socket, getWebSocket() replays the
    // existing subscriptions, and this one isn't among them yet. Sending it twice in the same
    // millisecond gives identical signatures, which the server rejects as a replay.
    this.subscriptions.add(resourceId);
  }

  /**
   * Unsubscribe from one resource topic.
   * @param resourceId - The resource ID to unsubscribe from
   */
  async unsubscribe(resourceId: string): Promise<void> {
    this.subscriptions.delete(resourceId);

    // Not connected: the server holds no topics for us, and future sockets won't replay this one.
    if (!this.webSocket) {
      return;
    }

    await this.request('/data/unsubscribe', { resourceId: [resourceId] });
  }

  /** Send a request and wait for the response carrying the same ID. */
  private async request(path: string, body: unknown): Promise<unknown> {
    const id = crypto.randomUUID();

    await this.send(path, body, id);

    // Registered after sending, so only requests that were actually sent can be rejected by a close.
    // The reply can't arrive first: socket messages are delivered as later events.
    return new Promise<unknown>((resolve, reject) => {
      this.pendingRequests.set(id, { resolve, reject });
    });
  }

  /**
   * Connect if needed, then sign and send one request envelope using the shared Extended JSON codec.
   * @param path - The path of the request
   * @param body - The body of the request
   * @param id - The ID of the request
   */
  private async send(path: string, body: unknown, id: string = crypto.randomUUID()): Promise<void> {
    const socket = await this.getWebSocket();

    // The socket can close between being awaited and being used. A closed socket drops frames silently.
    if (socket.readyState !== WebSocket.OPEN) {
      throw new Error('WebSocket is not connected');
    }

    // Signed after connecting: the server rejects signatures older than its timestamp window (5 minutes by default).
    const headers = BaseSyncClient.signRequest(this.privateKey.toBytes(), path, toExtendedJson(body));
    const request: WsRequest = { id, path, body, headers };

    socket.send(toExtendedJson(request));
  }

  /**
   * Resolve with the socket once it opens, or reject if the handshake fails.
   * @param socket - The WebSocket to wait for
   * @returns The WebSocket
   */
  private static waitForOpen(socket: WebSocket): Promise<WebSocket> {
    return new Promise((resolve, reject) => {
      socket.onopen = () => resolve(socket);
      socket.onerror = () => reject(new Error('WebSocket connection failed'));
    });
  }

  /**
   * Decode and route one frame received from the sync server.
   * @param raw - The raw message received from the sync server
   */
  private handleMessage(raw: string): void {
    const { success, data: message } = WsMessageSchema.safeParse(fromExtendedJson(raw));
    if (!success) {
      this.emitError(new Error('Invalid WebSocket message'));
      return;
    }

    // Only response/error messages with an outstanding ID are RPC replies.
    // Application events are left for message listeners, even if they have IDs.
    const pending = message.id ? this.pendingRequests.get(message.id) : undefined;

    // If the message has no outstanding ID, it is a subscription event so we will just notify the listeners
    if (!pending) {
      for (const listener of this.messageListeners) {
        listener(message);
      }

      // Errors with no pending request are failed subscribes (see sendSubscribe) or uncorrelated server errors.
      const error = WsErrorResponseSchema.safeParse(message);
      if (error.success) {
        this.emitError(new Error(`${error.data.error} (${error.data.statusCode})`));
        return;
      }

      if (message.type === 'instance-changed') {
        const parsedMessage = BaseSyncClient.resourceInstanceSchema.safeParse(message.data);
        if (!parsedMessage.success) {
          return;
        }

        this.emit('resource-instance-updated', {
          ...parsedMessage.data,
          value: fromExtendedJson(new TextDecoder().decode(parsedMessage.data.blob)) as PayloadType,
        });
      }

      return;
    }

    // Delete the pending request
    this.pendingRequests.delete(message.id!);

    // If the message is an error, reject the pending request
    if (message.type === 'error') {
      const error = WsErrorResponseSchema.parse(message);
      pending.reject(new Error(`${error.error} (${error.statusCode})`));
      return;
    }

    // If the message is a response, resolve the pending request
    if (message.type === 'response') {
      const response = WsSuccessResponseSchema.parse(message);

      // If the response is not successful, reject the pending request
      if (response.statusCode < 200 || response.statusCode >= 300) {
        pending.reject(new Error(`Request failed (${response.statusCode})`));
        return;
      }

      // If the response is successful, resolve the pending request
      pending.resolve(response.body);
    }
  }

  /**
   * Helper to emit an error to all error listeners
   * @param error - The error to emit
   */
  private emitError(error: Error): void {
    for (const listener of this.errorListeners) {
      listener(error);
    }
  }

  /**
   * Helper for when the client disconnects
   * @param error - The error to reject the pending requests with
   */
  private rejectPendingRequests(error: Error): void {
    for (const pending of this.pendingRequests.values()) {
      pending.reject(error);
    }

    this.pendingRequests.clear();
  }

  /**
   * Helper to convert an HTTP URL to a WebSocket URL
   * @param httpUrl - The HTTP URL to convert
   * @returns The WebSocket URL
   */
  static httpToWsUrl(httpUrl: string): string {
    const url = new URL(httpUrl);

    if (url.protocol === 'http:') {
      url.protocol = 'ws:';
    } else if (url.protocol === 'https:') {
      url.protocol = 'wss:';
    } else if (url.protocol !== 'ws:' && url.protocol !== 'wss:') {
      throw new Error(`Unsupported sync server protocol: ${url.protocol}`);
    }

    // The constructor accepts either the server root or the complete `/ws` URL.
    const pathname = url.pathname.replace(/\/+$/, '');
    url.pathname = pathname.endsWith('/ws') ? pathname : `${pathname}/ws`;

    return url.toString();
  }
}