sdk-signer-client/src/client.ts

import { WebSocket } from 'ws';
import { Data } from 'ws';
import {
  ClientConfig,
  ClientMessage,
  ServerResponse,
  ClientEvents,
  MessageType,
  ConnectionOptions,
  MessageHandler,
  ConnectionHandler,
  ErrorHandler
} from './types';

export class SDKSignerClient {
  private ws: WebSocket | null = null;
  private config: ClientConfig;
  private events: ClientEvents = {};
  private reconnectAttempts = 0;
  private reconnectTimer: NodeJS.Timeout | null = null;
  private messageHandlers = new Map<string, MessageHandler>();
  private isConnecting = false;
  private isConnected = false;

  constructor(config: ClientConfig) {
    this.config = {
      timeout: 5000,
      reconnectInterval: 3000,
      maxReconnectAttempts: 5,
      ...config
    };
  }

  /**
   * Connect to the SDK Signer server
   */
  async connect(options: ConnectionOptions = {}): Promise<void> {
    if (this.isConnecting || this.isConnected) {
      return;
    }

    this.isConnecting = true;

    return new Promise((resolve, reject) => {
      const timeout = options.timeout || this.config.timeout || 5000;
      const timeoutId = setTimeout(() => {
        this.isConnecting = false;
        reject(new Error('Connection timeout'));
      }, timeout);

      try {
        this.ws = new WebSocket(this.config.url, {
          headers: options.headers
        });

        this.ws.on('open', () => {
          clearTimeout(timeoutId);
          this.isConnecting = false;
          this.isConnected = true;
          this.reconnectAttempts = 0;
          console.log('✅ Connected to SDK Signer server');
          
          if (this.events.open) {
            this.events.open();
          }
          resolve();
        });

        this.ws.on('message', (data: Data) => {
          try {
            const response: ServerResponse = JSON.parse(data.toString());
            this.handleMessage(response);
          } catch (error) {
            console.error('Failed to parse message:', error);
          }
        });

        this.ws.on('close', (code: number, reason: Buffer) => {
          this.isConnected = false;
          console.log(`🔌 Connection closed: ${code} - ${reason.toString()}`);
          
          if (this.events.close) {
            this.events.close();
          }

          // Attempt to reconnect if not manually closed
          if (code !== 1000) {
            this.scheduleReconnect();
          }
        });

        this.ws.on('error', (error: Error) => {
          clearTimeout(timeoutId);
          this.isConnecting = false;
          console.error('❌ WebSocket error:', error);
          
          if (this.events.error) {
            this.events.error(error);
          }
          reject(error);
        });

      } catch (error) {
        clearTimeout(timeoutId);
        this.isConnecting = false;
        reject(error);
      }
    });
  }

  /**
   * Disconnect from the server
   */
  disconnect(): void {
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
      this.reconnectTimer = null;
    }

    if (this.ws) {
      this.ws.close(1000, 'Client disconnect');
      this.ws = null;
    }

    this.isConnected = false;
    this.reconnectAttempts = 0;
  }

  /**
   * Send a message to the server
   */
  send(message: ClientMessage): void {
    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
      throw new Error('WebSocket is not connected');
    }

    // Add API key if not present
    if (!message.apiKey && this.config.apiKey) {
      message.apiKey = this.config.apiKey;
    }

    // Generate message ID if not present
    if (!message.messageId) {
      message.messageId = this.generateMessageId();
    }

    this.ws.send(JSON.stringify(message));
  }

  /**
   * Send a message and wait for a specific response
   */
  async sendAndWait(
    message: ClientMessage,
    expectedType: MessageType,
    timeout: number = 10000
  ): Promise<ServerResponse> {
    return new Promise((resolve, reject) => {
      const timeoutId = setTimeout(() => {
        this.messageHandlers.delete(message.messageId!);
        reject(new Error(`Timeout waiting for ${expectedType} response to message ${message.type} (${message.messageId})`));
      }, timeout);

      const messageId = message.messageId || this.generateMessageId();
      message.messageId = messageId;

      const handler = (response: ServerResponse) => {
        if (response.messageId === messageId) {
          clearTimeout(timeoutId);
          this.messageHandlers.delete(messageId);
          
          if (response.type === MessageType.ERROR) {
            const errorMessage = response.error ? 
              (typeof response.error === 'string' ? response.error : JSON.stringify(response.error)) : 
              'Unknown server error';
            reject(new Error(`Server error for ${message.type}: ${errorMessage}`));
          } else if (response.type === expectedType) {
            resolve(response);
          } else {
            reject(new Error(`Unexpected response type: ${response.type}, expected: ${expectedType}`));
          }
        }
      };

      this.messageHandlers.set(messageId, handler);
      this.send(message);
    });
  }

  /**
   * Set event handlers
   */
  on(event: keyof ClientEvents, handler: ConnectionHandler | ErrorHandler | MessageHandler): void {
    this.events[event] = handler as any;
  }

  /**
   * Remove event handler
   */
  off(event: keyof ClientEvents): void {
    delete this.events[event];
  }

  /**
   * Check if client is connected
   */
  isConnectedToServer(): boolean {
    return this.isConnected && this.ws?.readyState === WebSocket.OPEN;
  }

  /**
   * Wait for server to send LISTENING message (connection establishment)
   */
  async waitForListening(): Promise<ServerResponse> {
    return new Promise((resolve, reject) => {
      const timeout = setTimeout(() => {
        reject(new Error('Timeout waiting for server LISTENING message'));
      }, this.config.timeout || 10000);

      const handler = (response: ServerResponse) => {
        if (response.type === MessageType.LISTENING) {
          clearTimeout(timeout);
          this.off('message');
          resolve(response);
        }
      };

      this.on('message', handler);
    });
  }
  async get_owned_processes(): Promise<ServerResponse> {
    const message: ClientMessage = {
      type: MessageType.GET_MY_PROCESSES,
      messageId: this.generateMessageId()
    };

    const response = await this.sendAndWait(message, MessageType.PROCESSES_RETRIEVED);

    return response;
  }

  /**
   * Notify an update for a process
   */
  async notifyUpdate(processId: string, stateId: string): Promise<ServerResponse> {
    const message: ClientMessage = {
      type: MessageType.NOTIFY_UPDATE,
      processId,
      stateId,
      messageId: this.generateMessageId()
    };

    return this.sendAndWait(message, MessageType.UPDATE_NOTIFIED);
  }

  /**
   * Validate a state
   */
  async validateState(processId: string, stateId: string): Promise<ServerResponse> {
    const message: ClientMessage = {
      type: MessageType.VALIDATE_STATE,
      processId,
      stateId,
      messageId: this.generateMessageId()
    };

    return this.sendAndWait(message, MessageType.STATE_VALIDATED);
  }

  /**
   * Update a process
   */
  async updateProcess(processId: string, newData: { [label: string]: any }, privateFields: string[], roles: any | null): Promise<ServerResponse> {
    const message: ClientMessage = {
      type: MessageType.UPDATE_PROCESS,
      processId,
      newData,
      privateFields,
      roles,
      messageId: this.generateMessageId()
    };

    return this.sendAndWait(message, MessageType.PROCESS_UPDATED);
  }

  private handleMessage(response: ServerResponse): void {
    // Call general message handler
    if (this.events.message) {
      this.events.message(response);
    }

    // Call specific message handler if exists
    if (response.messageId) {
      const handler = this.messageHandlers.get(response.messageId);
      if (handler) {
        handler(response);
        this.messageHandlers.delete(response.messageId);
      }
    }

    // Handle error responses
    if (response.type === MessageType.ERROR) {
      const errorMessage = response.error ? 
        (typeof response.error === 'string' ? response.error : JSON.stringify(response.error)) : 
        'Unknown server error';
      
      console.error('Server error:', errorMessage);
      console.error('Full error response:', JSON.stringify(response, null, 2));
      
      if (this.events.error) {
        this.events.error(new Error(errorMessage));
      }
    }
  }

  private scheduleReconnect(): void {
    if (this.reconnectAttempts >= (this.config.maxReconnectAttempts || 5)) {
      console.error('Max reconnection attempts reached');
      return;
    }

    this.reconnectAttempts++;
    const delay = (this.config.reconnectInterval || 3000) * this.reconnectAttempts;

    console.log(`🔄 Attempting to reconnect in ${delay}ms (attempt ${this.reconnectAttempts})`);

    this.reconnectTimer = setTimeout(async () => {
      try {
        await this.connect();
        if (this.events.reconnect) {
          this.events.reconnect();
        }
      } catch (error) {
        console.error('Reconnection failed:', error);
        this.scheduleReconnect();
      }
    }, delay);
  }

  private generateMessageId(): string {
    return `msg_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
  }
}