import * as ws from 'ws';
import {socketKey} from '../../../utils.js';
import {relays} from '../relays/core.js';
import getLogger from '../../../logger.js';
import {HttpServerRequest} from '../../../server/exchange.ts';

const logger = getLogger('mesh.ws.cluster');


function onMessage(key: string, node: string | undefined, socketsByNodeId: Map<string, Map<string, ws.WebSocket>>, msg: string) {
    try {
        relays.send(key, node, msg, (k, err?: Error) => {
            if (err) {
                logger.warn(`${k} error writing msg ${msg}: ${err}`);
                return;
            }
            if (logger.enabledFor('debug')) {
                logger.debug(`${k} sent msg ${msg}`);
            }
        });
    } catch (ex) {
        logger.error(`${key} unable to process message`, ex);
        if (node) {
            const socket = socketsByNodeId.get(node)?.get(key);
            socket?.terminate();
        }
    }
}

async function create(server: {wss: ws.WebSocketServer}): Promise<{ close?: () => Promise<void> }> {
    const socketsByNodeId = new Map<string, Map<string, ws.WebSocket>>();
    relays.onMsg = (k, nodeId, msg) => {
        try {
            const sockets = socketsByNodeId.get(nodeId);
            if (sockets && sockets.size > 0) {
                for (const [key, socket] of sockets) {
                    socket.send(msg, {binary: false}, (err?: Error) => {
                        if (err) {
                            logger.warn(`${key} error writing from ${k} msg ${msg}: ${err}`);
                            return;
                        }
                        if (logger.enabledFor('debug')) {
                            logger.debug(`${key} sent from ${k} msg ${msg}`);
                        }
                    });
                }
            } else {
                logger.warn(`${k} dropped msg ${msg}.`);
            }
        } catch (ex) {
            logger.error(`${k} unable to process message`, ex);
        }
    };
    server.wss
        .on('error', () => {
            logger.error(`error starting mesh server`);
        })
        .on('listening', () => {
            logger.info(`mesh server is listening`);
        })
        .on('connection', (socket, req) => {
            const request = new HttpServerRequest(req);
            const key = socketKey(request.socket);
            const query = new URLSearchParams(request.query ?? undefined);
            logger.info(`${key} connected on cluster with ${query}`);
            const node = query.get('node');
            if (node) {
                let sockets = socketsByNodeId.get(node);
                if (!sockets) {
                    sockets = new Map<string, ws.WebSocket>();
                    socketsByNodeId.set(node, sockets);
                }
                sockets.set(key, socket);
            }
            else {
                socket.terminate();
                return;
            }
            //onOpen(connectedNodes, key);
            socket.on('error', (err: Error) => {
                logger.error(`${key} websocket error: ${err}`, err);
            });
            socket.on('message', (data, _isBinary) => {
                if (Array.isArray(data)) {
                    data = Buffer.concat(data);
                }
                onMessage(key, node, socketsByNodeId,  data as unknown as string);
            });
            socket.on('close', (_code, _reason) => {
                logger.info(`${key} disconnected from cluster`);
                const sockets = socketsByNodeId.get(node);
                if (sockets) {
                    sockets.delete(key);
                    if (sockets.size === 0) {
                        socketsByNodeId.delete(node);
                    }
                }
            });
        });
    return {
        close: async () => {
            server.wss.close();
        }
    };
}

export default create;