import { WebSocket, WebSocketServer as WebSocketServerRaw } from 'ws';
import { Logger, red } from '../index.js';
const HMR_HEADER = 'farm_hmr';
const wsServerEvents = [
    'connection',
    'error',
    'headers',
    'listening',
    'message'
];
export default class WsServer {
    constructor(httpServer, config, hmrEngine, logger) {
        this.httpServer = httpServer;
        this.config = config;
        this.hmrEngine = hmrEngine;
        this.customListeners = new Map();
        this.clientsMap = new WeakMap();
        this.bufferedError = null;
        this.logger = logger ?? new Logger();
        this.createWebSocketServer();
    }
    createWebSocketServer() {
        try {
            this.wss = new WebSocketServerRaw({ noServer: true });
            this.connection();
            // TODO IF not have httpServer
            this.httpServer.on('upgrade', this.upgradeWsServer.bind(this));
        }
        catch (err) {
            this.handleSocketError(err);
        }
    }
    upgradeWsServer(request, socket, head) {
        if (this.isHMRRequest(request)) {
            this.handleHMRUpgrade(request, socket, head);
        }
    }
    listen() {
        // TODO alone with use httpServer we need start this function
        // Start listening for WebSocket connections
    }
    // Farm uses the `sendMessage` method in hmr and
    // the send method is reserved for migration vite
    send(...args) {
        let payload;
        if (typeof args[0] === 'string') {
            payload = {
                type: 'custom',
                event: args[0],
                data: args[1]
            };
        }
        else {
            payload = args[0];
        }
        if (payload.type === 'error' && !this.wss.clients.size) {
            this.bufferedError = payload;
            return;
        }
        const stringified = JSON.stringify(payload);
        this.wss.clients.forEach((client) => {
            // readyState 1 means the connection is open
            if (client.readyState === 1) {
                client.send(stringified);
            }
        });
    }
    isHMRRequest(request) {
        return (request.url === this.config.hmr.path &&
            request.headers['sec-websocket-protocol'] === HMR_HEADER);
    }
    handleHMRUpgrade(request, socket, head) {
        this.wss.handleUpgrade(request, socket, head, (ws) => {
            this.wss.emit('connection', ws, request);
        });
    }
    get clients() {
        return new Set(Array.from(this.wss.clients).map(this.getSocketClient.bind(this)));
    }
    // a custom method defined by farm to send custom events
    sendCustomEvent(event, payload) {
        // Send a custom event to all clients
        this.send({ type: 'custom', event, data: payload });
    }
    on(event, listener) {
        if (wsServerEvents.includes(event)) {
            this.wss.on(event, listener);
        }
        else {
            this.addCustomEventListener(event, listener);
        }
    }
    off(event, listener) {
        if (wsServerEvents.includes(event)) {
            this.wss.off(event, listener);
        }
        else {
            this.removeCustomEventListener(event, listener);
        }
    }
    connection() {
        this.wss.on('connection', (socket) => {
            socket.on('message', (raw) => {
                if (!this.customListeners.size)
                    return;
                let parsed;
                try {
                    parsed = JSON.parse(String(raw));
                }
                catch {
                    this.logger.error('Failed to parse WebSocket message: ' + raw);
                }
                // transform vite js-update to farm update
                if (parsed?.type === 'js-update' && parsed?.path) {
                    this.hmrEngine.hmrUpdate(parsed.path, true);
                    return;
                }
                if (!parsed || parsed.type !== 'custom' || !parsed.event)
                    return;
                const listeners = this.customListeners.get(parsed.event);
                if (!listeners?.size)
                    return;
                const client = this.getSocketClient(socket);
                listeners.forEach((listener) => listener(parsed.data, client));
            });
            socket.on('error', (err) => {
                return this.handleSocketError(err);
            });
            socket.send(JSON.stringify({ type: 'connected' }));
            if (this.bufferedError) {
                socket.send(JSON.stringify(this.bufferedError));
                this.bufferedError = null;
            }
        });
    }
    async close() {
        if (this.upgradeWsServer && this.httpServer) {
            this.httpServer.off('upgrade', this.upgradeWsServer);
        }
        await this.terminateAllClients();
        await this.closeWebSocketServer();
        // TODO if not have httpServer we need close httpServer
    }
    terminateAllClients() {
        const terminatePromises = Array.from(this.wss.clients).map((client) => {
            return new Promise((resolve) => {
                if (client.readyState === WebSocket.OPEN) {
                    client.send(JSON.stringify({ type: 'closing' }));
                    client.close(1000, 'Server shutdown');
                }
                // Temporarily remove the direct shutdown of ws
                // client.terminate();
                client.once('close', () => resolve(true));
            });
        });
        return Promise.all(terminatePromises);
    }
    closeWebSocketServer() {
        return new Promise((resolve, reject) => {
            this.wss.close((err) => {
                if (err) {
                    reject(err);
                }
                else {
                    // TODO if not have httpServer
                    resolve(true);
                }
            });
        });
    }
    addCustomEventListener(event, listener) {
        if (!this.customListeners.has(event)) {
            this.customListeners.set(event, new Set());
        }
        this.customListeners.get(event).add(listener);
    }
    removeCustomEventListener(event, listener) {
        this.customListeners.get(event)?.delete(listener);
    }
    getSocketClient(socket) {
        if (!this.clientsMap.has(socket)) {
            this.clientsMap.set(socket, {
                send: (...args) => this.sendMessage(socket, ...args),
                socket,
                rawSend: (str) => socket.send(str)
            });
        }
        return this.clientsMap.get(socket);
    }
    sendMessage(socket, ...args) {
        let payload;
        if (typeof args[0] === 'string') {
            payload = {
                type: 'custom',
                event: args[0],
                data: args[1]
            };
        }
        else {
            payload = args[0];
        }
        socket.send(JSON.stringify(payload));
    }
    handleSocketError(err) {
        if (err.code === 'EADDRINUSE') {
            this.logger.error(red(`WebSocket server error: Port is already in use`), {
                error: err
            });
        }
        else {
            this.logger.error(red(`WebSocket server error:\n${err.stack || err.message}`), {
                error: err
            });
        }
    }
}
//# sourceMappingURL=ws.js.map