import { PersonaPacket, PersonaProtocolBaseConfig, ProtocolStatus, Session } from '../types';
import { PersonaProtocolBase } from './base';

type PersonaWebSocketProtocolConfig = PersonaProtocolBaseConfig & {
  webSocketUrl: string;
};

class PersonaWebSocketProtocol extends PersonaProtocolBase {
  status: ProtocolStatus;
  autostart: boolean;
  session: Session;
  config: PersonaWebSocketProtocolConfig;
  webSocket: WebSocket | null;

  constructor(config: PersonaWebSocketProtocolConfig) {
    super();
    this.config = config;
    this.status = 'disconnected';
    this.autostart = true;
    this.session = null;
    this.webSocket = null;
  }

  public getName(): string {
    return 'websocket';
  }

  public getPriority(): number {
    return 500;
  }

  public async syncSession(session: Session): Promise<void> {
    this.config.logger?.debug('Syncing session with WebSocket protocol:', session);
    this.session = session;
    if (this.webSocket && this.status === 'connected') {
      this.disconnect();
      this.connect(session);
    }
  }

  public connect(session?: Session): Promise<Session> {
    if (this.webSocket !== null && this.status === 'connected') {
      return Promise.resolve(this.session);
    }

    const sid = session || this.session || 'new';

    this.config.logger?.debug('Connecting to WebSocket with sessionId:', sid);

    const apiKey = encodeURIComponent(this.config.apiKey);
    const agentId = this.config.agentId;
    const userId = this.config.userId || 'anonymous';
    const webSocketUrl = `${this.config.webSocketUrl}?sessionCode=${sid}&agentId=${agentId}&apiKey=${apiKey}&userId=${encodeURIComponent(
      userId,
    )}`;
    this.setStatus('connecting');
    this.webSocket = new WebSocket(webSocketUrl);
    this.webSocket.addEventListener('open', () => {
      this.setStatus('connected');
    });
    this.webSocket.addEventListener('message', (event) => {
      const data = JSON.parse(event.data) as PersonaPacket;
      this.notifyPacket(data);
    });
    this.webSocket.addEventListener('close', (event: CloseEvent) => {
      this.setStatus('disconnected');
      this.webSocket = null;
      if (event.code !== 1000) {
        this.notifyPacket({
          type: 'message',
          payload: {
            role: 'assistant',
            type: 'text',
            text: 'Oops! The connection to the server was lost. Please try again later.',
          },
        });
        this.config.logger?.warn('WebSocket connection closed');
      }
    });

    this.webSocket.addEventListener('error', () => {
      this.setStatus('disconnected');
      this.webSocket = null;
      this.config.logger?.error('WebSocket connection error');
    });

    return Promise.resolve(sid);
  }

  public disconnect(): Promise<void> {
    this.config.logger?.debug('Disconnecting WebSocket');
    if (this.webSocket && this.status === 'connected') {
      this.webSocket.close(1000, 'Normal closure');
      this.setStatus('disconnected');
      this.webSocket = null;
    }
    return Promise.resolve();
  }

  public sendPacket(packet: PersonaPacket): Promise<void> {
    if (this.webSocket && this.status === 'connected') {
      this.webSocket.send(JSON.stringify(packet));
      return Promise.resolve();
    } else {
      return Promise.reject(new Error('WebSocket is not connected'));
    }
  }
}

export { PersonaWebSocketProtocol };
export type { PersonaWebSocketProtocolConfig };