import WebSocket from 'ws';

export interface GradioClientOptions {
  timeout?: number;
}

// eslint-disable-next-line @typescript-eslint/ban-types
export class GradioClient<M extends Record<number, [any[], any[]]> = {}> {
  constructor(
    public endpoint: string,
    private options: GradioClientOptions = {},
  ) {}

  async request<I extends number & keyof M>(
    i: I,
    ...args: M[I][0]
  ): Promise<M[I][1]> {
    const ws = new WebSocket(`${this.endpoint}/queue/join`);
    const timeout = this.options.timeout || 60000;
    const jobId = Math.floor(Math.random() * Number.MAX_SAFE_INTEGER).toString(
      36,
    );
    return new Promise((resolve, reject) => {
      let resolved = false;
      const tmot = setTimeout(() => {
        ws.close();
        reject(new Error(`Timeout of ${timeout}ms exceeded`));
      }, timeout);
      ws.onopen = () => {
        // ws.send(JSON.stringify({ i, args }));
      };
      ws.onmessage = (e) => {
        const data: { msg: string; output?: { data: any[] } } = JSON.parse(
          e.data as string,
        );
        const base = { fn_index: i, session_hash: jobId };
        switch (data.msg) {
          case 'send_hash':
            ws.send(JSON.stringify(base));
            break;
          case 'send_data':
            ws.send(JSON.stringify({ ...base, event_data: null, data: args }));
            break;
          case 'process_completed':
            clearTimeout(tmot);
            resolved = true;
            ws.close();
            resolve(data.output.data);
            break;
        }
      };
      ws.onerror = (e) => {
        clearTimeout(tmot);
        if (!resolved) {
          reject(new Error(`WebSocket error: ${e.message}`));
        }
      };
      ws.onclose = (e) => {
        clearTimeout(tmot);
        if (!resolved) {
          reject(
            new Error(`WebSocket closed unexpectedly: ${e.code} ${e.reason}`),
          );
        }
      };
    });
  }
}