const { gpt, gptweb } = require("./API/gpt");
const { dalle } = require("./API/dalle");
const { gpt4o } = require("./API/gpt4o");
const { gpt3 } = require("./API/gpt3.5");
const { prompts } = require("./API/gpt-prompts");
const { blackbox } = require("./API/blackbox");
const { qwen } = require("./API/qwen");
const { gemini } = require("./API/gemini");
const WebSocket = require('ws');
const Bottleneck = require('bottleneck');
const { stats, scheduleDailyReset, resetStatistics } = require('./stats'); // Импортируем модуль статистики
const fs = require('fs');
const path = require('path');
const Ajv = require('ajv');

// Создаем экземпляр Ajv (проверка входящей переписки на соответствие схеме)
const ajv = new Ajv(); 
// Объект для хранения действий и состояния каждого клиента
const clientActions = {};
// Объект для хранения ограничителей по каждому клиенту
const limiters = {}; 
// Замените на ваш домен
const allowedOrigin = 'http://gpt.loc'; 

// Создайте WebSocket-сервер
const wss = new WebSocket.Server({ port: 3000 });

wss.on('connection', (ws, req) => {
	
	const urlParams = new URLSearchParams(req.url.split('?')[1]);
    const userId = urlParams.get('token'); // Получаем UUID из URL
    const userIP = req.socket.remoteAddress;
	//console.log("IP-адрес - ", userIP, "Время подключения - ", new Date().toISOString());
	const userHeaders = req.headers;

	// Проверяем Origin
    if (!checkOrigin(userHeaders.origin, ws)) {
        return; // Если отклонено, выходим из обработчика
    }
	
	// Увеличиваем общий счетчик подключений
    stats.totalConnections++;
	
	console.log('Client connected with user ID:', userId, 'Домен - ', userHeaders.origin);
	
	// Инициализируем состояние клиента, если его еще нет
    if (!clientActions[userId]) {
		clientActions[userId] = {
			actions: [], // Массив для хранения действий
			connected: true, // Статус подключения
        };
	}
	else
		clientActions[userId].connected = true;
	
	if (!limiters[userId]) { 
	  // Создаем экземпляр Bottleneck для каждого клиента (ограничение на количество запросов по Id)
		limiters[userId] = new Bottleneck({
            maxConcurrent: 1, // Максимальное количество одновременных задач
            minTime: 1000, // Минимальное время между запросами в миллисекундах
		    highWater: 0 // Не разрешать в очереди больше 0 запросов (игнорировать)
        });
    }

    ws.on('message', (message) => {
		// Проверяем, является ли сообщение экземпляром Buffer
		if (Buffer.isBuffer(message))
			message = message.toString(); 
		try {
			// Преобразуем строку JSON обратно в объект
			message = JSON.parse(message); 
			//Получаем историю чата из запроса
			let messages = message.messages;
			let model = message.model;
			console.log('Модель - ', model);
			// Выполняем валидацию сообщений
			messages = validateMessages(messages);
			
			// Проверяем, остались ли сообщения после валидации
			if (messages.length === 0) {
				const response = { code: 2, message: "Нет валидных сообщений для обработки." };
				sendMessage(ws, response);
				return;
			}
			
			// Урезаем последнее сообщение до 2000 символов, если оно от клиента
			if (messages[messages.length - 1].role === 'user') {
				messages[messages.length - 1].content = messages[messages.length - 1].content.substring(0, 2000000);
			}
			
			// Увеличиваем общий счетчик запросов
            stats.totalRequests++;
			
			// Используем Bottleneck для управления количеством запросов
			limiters[userId].schedule(() => {
				// Отправка запроса на сервер чата
				if (model == 'gpt-4o')
					return callGPT4o(messages, ws);
				else if (model == 'gemini-pro')
					return callGemini(messages, ws);
				else if (model == 'qwen')
					return callQwen(messages, ws);
				else if (model == 'blackbox')
					return callBlackbox(messages, ws);
				else if (model == 'gpt-3.5-turbo')
					return callGPT35Turbo(messages, ws);
			}).catch(error => {
				stats.totalIPBlocks++;
				//console.error('Rate limit exceeded, message ignored:', error);
				const response = { code: 2, message: "Превышен лимит запросов. Пожалуйста, попробуйте позже." };
				// Отправка ответа клиенту
				//sendMessage(ws, response);
			});
			
		} catch (error) {
				console.error('Error parsing JSON:', error);
				const response = { code: 2, message: "Произошла ошибка. Попробуйте еще раз."};
				console.log(response);
				// Отправка ответа клиенту
				sendMessage(ws, response);
		}
		
    });

    ws.on('close', () => {
		clientActions[userId].connected = false; // Обновление статуса подключения
		delete clientActions[userId]; // Удаление состояния клиента
		delete limiters[userId]; // Удаление ограничителя для данного клиента
        console.log('Client ',userId, ' disconnected');
    });
});

// Запускаем планировщик для статистики
scheduleDailyReset();

// Очистка старых соединений раз в минуту (по желанию)
setInterval(() => {
  Object.keys(clientActions).forEach(userId => {
    if (!clientActions[userId].connected) {
      delete clientActions[userId]; // Удаляем пользователя из памяти, если он отключен
	  delete limiters[userId];
    }
  });
}, 10 * 60 * 1000); // Очистка каждые 10 минут


function checkOrigin(origin, socket) {
    if (origin !== allowedOrigin) {
        console.log('Недопустимый домен. Соединение разорвано: ', origin);
        socket.close(); // Закрываем соединение, если домен не разрешен
        return false; // Возвращаем false для индикатора отклонения
    }
    return true; // Возвращаем true, если домен разрешен
}

// Отправляет ответ клиенту
function sendMessage(ws, message) {
    if (ws.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify(message));
    } else {
        console.warn(`Trying to send message to a closed WebSocket connection.`);
    }
}

// Удаление из массива сообщений не соответствующих схеме
function validateMessages(messages) {
    // Определяем схему для сообщений
    const schema = {
        type: 'object',
        properties: {
            role: { type: 'string' },
            content: { type: 'string' }
        },
        required: ['role', 'content'],
        additionalProperties: false
    };

    // Фильтруем массив, оставляя только корректные сообщения
    return messages.filter(message => ajv.validate(schema, message));
}

console.log('WebSocket сервер запущен через порт 3000');

















//============== ВЫЗОВ API МОДЕЛЕЙ =================================================
//==================================================================================
//==================================================================================
//==================================================================================

// GPT4o
async function callGPT4o(messages, ws) {
    try {
        let fullResponse = ""; // Переменная для результата
		
		console.log('Запрос клиента:', messages[messages.length - 1]); 
		//throw new Error("Намеренная ошибка для тестирования catch");
        gpt4o({
            messages: messages,
            markdown: false
        }, (result) => {
			// Обработка результата внутри колбэка
            if (result.error) {
                throw new Error(result.error); // Обрабатываем ошибку, если она есть
            }
			try {
				if (!result.finish) {
					// Запоминаем только последний ответ чата (он будет полный, промежуточные ответы - это накопления)
					fullResponse = result.message || "";
					//console.log(result);
				} else {
					// Если завершено, выводим полный ответ
					let response = { code: 1, message: fullResponse};
					console.log('Ответ чата: ', response);
					// Отправка ответа клиенту
					sendMessage(ws, response);
				}
			}
			catch (callbackError) {
                console.error('Error processing callback result:', callbackError);
                let response = { code: 2, message: "Произошла ошибка. Попробуйте еще раз." };
                // Отправка ответа клиенту
				sendMessage(ws, response);
            }
        });
    } catch (error) {
        console.error('Error:', error);
		let response = { code: 2, message: "Произошла ошибка. Попробуйте еще раз."};
		// Отправка ответа клиенту
		sendMessage(ws, response);
    }
}

// GPT3.5-Turbo
async function callGPT35Turbo(messages, ws) {
    try {
        let fullResponse = ""; // Переменная для результата
		
		console.log('Запрос клиента:', messages[messages.length - 1]); 
		//throw new Error("Намеренная ошибка для тестирования catch");
        gpt3({
            messages: messages,
            markdown: false
        }, (result) => {
			// Обработка результата внутри колбэка
            if (result.error) {
                throw new Error(result.error); // Обрабатываем ошибку, если она есть
            }
			try {
				if (!result.finish) {
					// Запоминаем только последний ответ чата (он будет полный, промежуточные ответы - это накопления)
					fullResponse = result.message || "";
					//console.log(result);
				} else {
					// Если завершено, выводим полный ответ
					let response = { code: 1, message: fullResponse};
					console.log('Ответ чата: ', response);
					// Отправка ответа клиенту
					sendMessage(ws, response);
				}
			}
			catch (callbackError) {
                console.error('Error processing callback result:', callbackError);
                let response = { code: 2, message: "Произошла ошибка. Попробуйте еще раз." };
                // Отправка ответа клиенту
				sendMessage(ws, response);
            }
        });
    } catch (error) {
        console.error('Error:', error);
		let response = { code: 2, message: "Произошла ошибка. Попробуйте еще раз."};
		// Отправка ответа клиенту
		sendMessage(ws, response);
    }
}


// Gemini
async function callGemini(messages, ws) {
    try {
        let fullResponse = ""; // Переменная для результата
		
		console.log('Запрос клиента:', messages[messages.length - 1]); 
		//throw new Error("Намеренная ошибка для тестирования catch");
        gemini({
            messages: messages
        }, (result) => {
			// Обработка результата внутри колбэка
            if (result.error) {
                throw new Error(result.error); // Обрабатываем ошибку, если она есть
            }
			try {
				if (!result.finish) {
					// Запоминаем только последний ответ чата (он будет полный, промежуточные ответы - это накопления)
					fullResponse = result.message || "";
					//console.log(result);
				} else {
					// Если завершено, выводим полный ответ
					let response = { code: 1, message: fullResponse};
					console.log('Ответ чата: ', response);
					// Отправка ответа клиенту
					sendMessage(ws, response);
				}
			}
			catch (callbackError) {
                console.error('Error processing callback result:', callbackError);
                let response = { code: 2, message: "Произошла ошибка. Попробуйте еще раз." };
                // Отправка ответа клиенту
				sendMessage(ws, response);
            }
        });
    } catch (error) {
        console.error('Error:', error);
		let response = { code: 2, message: "Произошла ошибка. Попробуйте еще раз."};
		// Отправка ответа клиенту
		sendMessage(ws, response);
    }
}


// Qwen
async function callQwen(messages, ws) {
    try {
        let fullResponse = ""; // Переменная для результата
		
		console.log('Запрос клиента:', messages[messages.length - 1]); 
		//throw new Error("Намеренная ошибка для тестирования catch");
        qwen({
            messages: messages,
            markdown: false
        }, (result) => {
			// Обработка результата внутри колбэка
            if (result.error) {
                throw new Error(result.error); // Обрабатываем ошибку, если она есть
            }
			try {
				if (!result.finish) {
					// Запоминаем только последний ответ чата (он будет полный, промежуточные ответы - это накопления)
					fullResponse = result.message || "";
					//console.log(result);
				} else {
					// Если завершено, выводим полный ответ
					let response = { code: 1, message: fullResponse};
					console.log('Ответ чата: ', response);
					// Отправка ответа клиенту
					sendMessage(ws, response);
				}
			}
			catch (callbackError) {
                console.error('Error processing callback result:', callbackError);
                let response = { code: 2, message: "Произошла ошибка. Попробуйте еще раз." };
                // Отправка ответа клиенту
				sendMessage(ws, response);
            }
        });
    } catch (error) {
        console.error('Error:', error);
		let response = { code: 2, message: "Произошла ошибка. Попробуйте еще раз."};
		// Отправка ответа клиенту
		sendMessage(ws, response);
    }
}

// Blackbox
async function callBlackbox(messages, ws) {
    try {
        let fullResponse = ""; // Переменная для результата
		
		console.log('Запрос клиента:', messages[messages.length - 1]); 
		//throw new Error("Намеренная ошибка для тестирования catch");
        blackbox({
            messages: messages,
            markdown: false,
            websearch: false,
            llama: false
        }, (result) => {
			// Обработка результата внутри колбэка
            if (result.error) {
                throw new Error(result.error); // Обрабатываем ошибку, если она есть
            }
			try {
				if (!result.finish) {
					// Запоминаем только последний ответ чата (он будет полный, промежуточные ответы - это накопления)
					fullResponse = result.message || "";
					//console.log(result);
				} else {
					// Если завершено, выводим полный ответ
					let response = { code: 1, message: fullResponse};
					console.log('Ответ чата: ', response);
					// Отправка ответа клиенту
					sendMessage(ws, response);
				}
			}
			catch (callbackError) {
                console.error('Error processing callback result:', callbackError);
                let response = { code: 2, message: "Произошла ошибка. Попробуйте еще раз." };
                // Отправка ответа клиенту
				sendMessage(ws, response);
            }
        });
    } catch (error) {
        console.error('Error:', error);
		let response = { code: 2, message: "Произошла ошибка. Попробуйте еще раз."};
		// Отправка ответа клиенту
		sendMessage(ws, response);
    }
}

//============== ВЫЗОВ API МОДЕЛЕЙ (КОНЕЦ) ======================

На мой сайт большая нагрузка, много клиентов. Нужно добавить балансировку для распределения нагрузки между серверами. Что тут сделать?