import { JSONRPCClient, JSONRPCParams, JSONRPCServer } from 'json-rpc-2.0';
import WebSocket from 'ws';
import { nanoid } from 'nanoid';
import { ErrorInfo, TokenDetails, TokenParams, TokenRequest } from 'ably';

const webSocket = new WebSocket('ws://localhost:3000');

export const serverRpc = new JSONRPCServer();

const idToCallback: Record<string, any> = {};
const idToAuthCallback: Record<
  string,
  (tokenParams: TokenParams) => PromiseLike<string | TokenDetails | TokenRequest>
> = {};

serverRpc.addMethod('callback', ({ callbackId, args }: any) => {
  // console.log(`received callback JSON RPC request for callbackId=${callbackId}, with args=`, args);

  // TODO: need to handle non-serializable values here

  const cbArgs = Object.values(args);
  idToCallback[callbackId]?.(...cbArgs);
});

serverRpc.addMethod('authCallback', async ({ authCallbackId, tokenParams }: any) => {
  // console.log('auth callback: %s', authCallbackId);

  const result = await idToAuthCallback[authCallbackId](tokenParams);

  return {
    type: typeof result === 'string' ? 'string' : (result as TokenRequest).keyName ? 'TokenRequest' : 'TokenDetails',
    response: result,
  };
});

export const jsonRpc = new JSONRPCClient(
  (request) => {
    try {
      webSocket.send(JSON.stringify(request));
      return Promise.resolve();
    } catch (error) {
      return Promise.reject(error);
    }
  },
  () => nanoid(),
);

let webSocketReady = false;
let waiterForReadiness: () => void | undefined;

webSocket.on('message', (message) => {
  // console.log('received: %s', message);
  const data = JSON.parse(message.toString());

  if (data.jsonrpc && (data.result || data.error)) {
    jsonRpc.receive(data);
  } else if (data.status === 'ready') {
    webSocketReady = true;
    waiterForReadiness?.();
  } else if (data.jsonrpc && data.method) {
    serverRpc.receive(data).then((response) => {
      if (response?.result) webSocket.send(JSON.stringify(response));
    });
  }
});

webSocket.on('open', () => {
  webSocket.send(JSON.stringify({ role: 'UNIFIED_CLIENT' }));
});

export const waitForConnection = () =>
  new Promise<void>((resolve) => {
    if (webSocketReady) {
      resolve();
    } else {
      waiterForReadiness = resolve;
    }
  });

export const registerCallback = (id: string, cb: any) => (idToCallback[id] = cb);
export const registerAuthCallback = (
  id: string,
  cb: (tokenParams: TokenParams) => PromiseLike<string | TokenDetails | TokenRequest>,
) => (idToAuthCallback[id] = cb);

export const findCallback = (cb: any): string =>
  Object.entries(idToCallback).find(([id, callback]) => callback === cb)![0];

export class AsyncClient {
  async request(method: string, params: JSONRPCParams): Promise<any> {
    try {
      return await jsonRpc.request(method, params);
    } catch (e: any) {
      if (e.data?.ablyError) {
        const { code, message, statusCode } = e.data.errorInfo;
        throw new ErrorInfo(message, code, statusCode);
      } else {
        throw e;
      }
    }
  }
}