import WebSocket from 'ws';
import net from 'net';

const DEFAULT_REMOTE_SERVER = 'wss://ingress.e4mc.link';
const DEFAULT_TCP_ADDRESS = 'localhost';
const DEFAULT_TCP_PORT = 25565;

interface DomainAssignedMessage {
  DomainAssigned: string;
}

interface ChannelOpenMessage {
  ChannelOpen: [number, string];
}

interface ChannelClosedMessage {
  ChannelClosed: number;
}

type ServerMessage = DomainAssignedMessage | ChannelOpenMessage | ChannelClosedMessage;

class e4mcClient {
  private tcpAddress: string;
  private tcpPort: number;
  private channels: Map<number, net.Socket>;
  private pendingData: Map<number, Buffer[]>;
  private ws: WebSocket;
  private onDomainAssigned: ((domain: string) => void) | undefined;

  constructor(onDomainAssigned?: (domain: string) => void, tcpPort?: number, tcpAddress?: string, remoteServer?: string) {
    this.tcpAddress = tcpAddress ?? DEFAULT_TCP_ADDRESS;
    this.tcpPort = tcpPort ?? DEFAULT_TCP_PORT;
    this.channels = new Map();
    this.pendingData = new Map();
    this.onDomainAssigned = onDomainAssigned;

    this.ws = new WebSocket(remoteServer ?? DEFAULT_REMOTE_SERVER);
    this.ws.on('open', () => this.onOpen());
    this.ws.on('message', (data: Buffer, isBinary: boolean) => this.onMessage(data, isBinary));
  }

  onOpen() {
    console.log('Connected to the remote server.');
  }

  onMessage(data: Buffer, isBinary: boolean) {
    if (!isBinary) {
      const message: ServerMessage = JSON.parse(data.toString());

      if ('DomainAssigned' in message) {
        console.log('Domain assigned:', message.DomainAssigned);
        if (this.onDomainAssigned) {
          this.onDomainAssigned(message.DomainAssigned);
        }
      } else if ('ChannelOpen' in message) {
        this.onChannelOpen(...message.ChannelOpen);
      } else if ('ChannelClosed' in message) {
        this.onChannelClosed(message.ChannelClosed);
      }
    } else {
      this.onBinaryMessage(data);
    }
  }

  onChannelOpen(channelId: number, clientInfo: string) {
    console.log(`Opening channel ${channelId} for ${clientInfo}`);

    const tcpSocket = net.createConnection({ host: this.tcpAddress, port: this.tcpPort }, () => {
      console.log(`TCP connection established for channel ${channelId}`);
    });

    tcpSocket.on('data', (chunk) => {
      const data = Buffer.concat([Buffer.from([channelId]), chunk]);
      this.ws.send(data);
    });

    tcpSocket.on('close', () => {
      console.log(`TCP connection closed for channel ${channelId}`);
      this.ws.send(JSON.stringify({ ChannelClosed: channelId }));
      this.channels.delete(channelId);
    });

    this.channels.set(channelId, tcpSocket);

    if (this.pendingData.has(channelId)) {
      this.pendingData.get(channelId)?.forEach((chunk) => {
        tcpSocket.write(chunk);
      });
      this.pendingData.delete(channelId);
    }
  }

  onChannelClosed(channelId: number) {
    console.log(`Closing channel ${channelId}`);
    if (this.channels.has(channelId)) {
      this.channels.get(channelId)?.end();
      this.channels.delete(channelId);
    }
  }

  onBinaryMessage(data: Buffer) {
    const channelId = data.readUInt8(0);
    const payload = data.slice(1);
    if (this.channels.has(channelId)) {
      this.channels.get(channelId)?.write(payload);
    } else {
      if (!this.pendingData.has(channelId)) {
        this.pendingData.set(channelId, []);
      }
      this.pendingData.get(channelId)?.push(payload);
    }
  }
}

export default e4mcClient;