import express, { json, urlencoded } from 'express';
import cors from 'cors';
import dotenv from 'dotenv';
import WebSocket, { WebSocketServer } from 'ws';
import { Worker } from 'worker_threads';
import { checkUsertoken } from './school/api.js';
import redis from './redis/index.js';
import initQueue from './redis/initQueue.js';
import messagePersistence from './redis/messagePersistence.js';
import code from './config/code.json' with { type: 'json' };
import fileRouter from './upload/index.js';
import recordRouter from './outside/callback.js';
import mdWebSocketServer from './utils/mdWebSocketServer.js';

dotenv.config();

const logger = {
  info: (message) => {
    const timestamp = new Date().toISOString();
    console.log(`[${timestamp}] INFO: ${message}`);
  },
  error: (message, error) => {
    const timestamp = new Date().toISOString();
    console.error(`[${timestamp}] ERROR: ${message}`, error || '');
  },
  debug: (message) => {
    const timestamp = new Date().toISOString();
    console.debug(`[${timestamp}] DEBUG: ${message}`);
  }
};

let wss = null;
const workers = [];
const socketMap = new Map();

function createWorker(scriptPath) {
  const worker = new Worker(scriptPath);
  worker.setMaxListeners(20);
  worker.on('error', (error) => {
    logger.error(`Worker ${scriptPath} 错误:`, error);
  });
  worker.on('exit', (code) => {
    if (code !== 0) {
      logger.error(`Worker ${scriptPath} 异常退出，退出码: ${code}`);
    }
  });
  workers.push(worker);
  return worker;
}

const assessment = createWorker('./worker_threads/assessment/assessment.js');
const wait = createWorker('./worker_threads/wait/waiting.js');
const polling = createWorker('./worker_threads/process/process.js');
const result = createWorker('./worker_threads/result/result.js');
const callback_timeout = createWorker('./worker_threads/callback_timeout/callbackTimeout.js');
const callback_result = createWorker('./worker_threads/callback_result/result.js');
const error = createWorker('./worker_threads/error/error.js');

async function sendMessageToClient(id, message, close = false, closeCode = 1000, closeReason = '') {
  let socket;
  if (typeof id === 'string' && id) {
    socket = socketMap.get(id);
  }
  if (socket && socket.readyState === WebSocket.OPEN && message) {
    try {
      socket.send(message);
      const messagePreview = typeof message === 'string' ? message.slice(0, 50) : JSON.stringify(message).slice(0, 50);
      logger.debug(`成功发送消息到客户端，id: ${id}, 消息: ${messagePreview}...`);
      if (close) {
        socket.close(closeCode, closeReason);
      }
      return true;
    } catch (error) {
      logger.error(`发送消息给客户端失败，id: ${id}`, error);
      return false;
    }
  } else {
    if (!message) {
      logger.debug(`消息为空，无法发送，id: ${id}`);
      return false;
    } else {
      logger.debug(`未找到目标客户端或连接已关闭，保存消息到待发送队列，id: ${id}`);
      try {
        await messagePersistence.savePendingMessage(id, message);
        logger.info(`消息已保存到待发送队列，等待重试: backendId=${id}`);
        return false;
      } catch (error) {
        logger.error(`保存待发送消息失败: backendId=${id}`, error);
        return false;
      }
    }
  }
}

function createWebSocketServer(server) {
  wss = new WebSocketServer({ 
    server,
    verifyClient: async (info, callback) => {
      try {
        const urlParams = new URLSearchParams(info.req.url.split('?')[1]);
        const token = urlParams.get('token');
        const id = urlParams.get('id');
        
        if (!token) {
          logger.info('缺少令牌');
          callback(false, 401, '缺少令牌');
          return;
        } else if (token !== process.env.TOKEN_SECRET){
          logger.info('验证后端失败');
          callback(false, 401, 'Token is invalid');
          return;
        }
        
        info.req.id = id;
        logger.info(`用户ID: token 验证成功`);
        callback(true);
      } catch (error) {
        logger.error('验证后端失败:', error);
        callback(false, 401, 'Token is invalid');
      }
    }
  });
  
  logger.info(`WebSocket server is running on the same port as HTTP server`);
  
  wss.on('error', (error) => {
    logger.error('WebSocket服务器错误:', error);
  });
  
  wss.on('connection', async (socket, req) => {
    const id = req.id;
    logger.info(`${id}号后端 连接成功`);
    socketMap.set(id, socket);
    socket.send('please give me tasks');

    try {
      const pendingMessages = await messagePersistence.getPendingMessages(id);
      if (pendingMessages.length > 0) {
        logger.info(`${id}号后端 发现 ${pendingMessages.length} 条待发送消息，开始重试发送`);
        
        for (const pendingMsg of pendingMessages) {
          try {
            socket.send(pendingMsg.message);
            await messagePersistence.removePendingMessage(pendingMsg.key);
            logger.debug(`成功重试发送消息: backendId=${id}, messageKey=${pendingMsg.key}`);
          } catch (error) {
            logger.error(`重试发送消息失败: backendId=${id}, messageKey=${pendingMsg.key}`, error);
            await messagePersistence.incrementRetryCount(pendingMsg.key);
          }
        }
        
        logger.info(`${id}号后端 待发送消息重试完成`);
      }
    } catch (error) {
      logger.error(`获取或发送待发送消息失败: backendId=${id}`, error);
    }

    socket.on('message', (message) => {
      const messageStr = typeof message === 'string' ? message : message.toString();
      
      if (messageStr === 'ping') {
        socket.send('pong');
        return;
      }
      
      if (messageStr === 'pong') {
        return;
      }
      
      if (messageStr === 'please give me tasks') {
        return;
      }

      try {
        const msg = JSON.parse(messageStr);
        
        if (msg.type === 'JWT_UPDATE') {
          logger.info(`收到 JWT_UPDATE 消息`);
          return;
        }
        
        if (msg.type === 'CAPACITY_UPDATE') {
          logger.debug(`收到算力状态更新: 可用容量 = ${msg.data?.summary?.availableCapacity || 0}`);
          return;
        }
        
        if (msg.type === 'INSTANCE_ONLINE') {
          logger.debug(`收到实例上线: ${msg.data?.instanceId}`);
          return;
        }
        
        if (msg.type === 'INSTANCE_OFFLINE') {
          logger.debug(`收到实例下线: ${msg.data?.instanceId}`);
          return;
        }
        
        if (msg.type === 'HEARTBEAT') {
          socket.send(JSON.stringify({
            type: 'HEARTBEAT_ACK',
            data: { timestamp: new Date().toISOString() }
          }));
          return;
        }
        
        const prefix = messageStr.slice(0, 50);
        if (prefix.includes('"type":"generate"') || prefix.includes("'type':'generate'")) {
          assessment.postMessage({
            type: 'submit',
            data: messageStr
          });
        } else {
          logger.debug(`收到未知消息类型: ${prefix}`);
        }
      } catch (e) {
        logger.error('处理消息出错:', e);
        socket.send(JSON.stringify({ 
          error: '处理消息出错',
          details: e.message 
        }));
      }
    });
    
    const heartbeatInterval = setInterval(() => {
      if (socket.readyState === WebSocket.OPEN) {
        socket.send('ping');
        // logger.debug(`向 ${id} 号后端发送心跳`);
      }
    }, 30000);
    
    socket.on('close', (code, reason) => {
      clearInterval(heartbeatInterval);
      logger.info(`${id}号后端 连接关闭，关闭码: ${code}，原因: ${reason}`);
    });
    
    socket.on('error', (error) => {
      logger.error(`${id}号后端 连接错误:`, error);
    });
  });
  
  assessment.on('message', async (message) => {
    logger.debug(`收到assessment worker消息: ${JSON.stringify(message)}`);
    if (message.type === 'AssessmentSuccess') {
      await sendMessageToClient(message.backendId, code.SUCCESS[message.type]);
    } else {
      await sendMessageToClient(message.backendId, code.ERROR[message.type], false, 4401, code.ERROR[message.type]);
    }
  });
  
  result.on('message', async (message) => {
    logger.debug(`收到result worker消息: ${JSON.stringify(message)}`);
    if (message.type === 'success') {
      await sendMessageToClient(message.backendId, message.message, false, 1000, 'success');
    } else {
      await sendMessageToClient(message.backendId, '获取结果失败，可在历史记录区刷新查看结果', false, 4401, code.ERROR[message.type]);
    }
  });
  
  callback_result.on('message', async (message) => {
    logger.debug(`收到callback_result worker消息: ${JSON.stringify(message)}`);
    if (message.type === 'success') {
      await sendMessageToClient(message.backendId, message.message, false, 1000, 'success');
    } else {
      await sendMessageToClient(message.backendId, '获取结果失败，可在历史记录区刷新查看结果', false, 4401);
    }
  });
  
  error.on('message', async (message) => {
    logger.debug(`收到error worker消息: ${JSON.stringify(message)}`);
    await sendMessageToClient(message.backendId, message.message, false, 4402, 'false');
  });
}

function gracefulShutdown() {
  logger.info('开始优雅关闭...');
  
  if (wss) {
    wss.close(() => {
      logger.info('WebSocket服务器已关闭');
    });
    
    wss.clients.forEach((client) => {
      client.close(1001, '服务器正在关闭');
    });
  }
  
  workers.forEach((worker, index) => {
    logger.info(`终止worker线程 ${index}`);
    worker.terminate();
  });
  
  if (redis.isOpen) {
    redis.disconnect()
      .then(() => {
        logger.info('Redis连接已关闭');
        process.exit(0);
      })
      .catch((error) => {
        logger.error('关闭Redis连接失败:', error);
        process.exit(1);
      });
  } else {
    process.exit(0);
  }
  
  setTimeout(() => {
    logger.info('强制关闭...');
    process.exit(1);
  }, 5000);
}

async function initialize() {
  logger.info('***************初始化队列开始***************');
  try {
    if (!redis.isOpen) {
      await redis.connect();
      logger.info('Redis 连接成功');
    }
    
    await initQueue.init();
    logger.info('***************初始化队列完成***************');
    
    setInterval(async () => {
      try {
        await messagePersistence.cleanupOldMessages(2 * 24 * 60 * 60 * 1000);
      } catch (error) {
        logger.error('定期清理过期消息失败:', error);
      }
    }, 24 * 60 * 60 * 1000);
    
  } catch (err) {
    logger.error('初始化失败:', err);
    process.exit(1);
  }
}

const app = express();
const hostname = '0.0.0.0';
const port = process.env.WS_PORT || 8087;

app.use(cors());
app.use('/workflow/uploads', express.static('uploads'));
app.use(json());
app.use(urlencoded({ extended: true }));
app.use('/workflow/file', fileRouter);
app.use('/callback', recordRouter);

app.get('/', (req, res) => {
  res.json({
    name: 'ComfyUI Task Queue Backend',
    version: '1.0.0',
    status: 'running',
    timestamp: new Date().toISOString()
  });
});

process.on('SIGINT', gracefulShutdown);
process.on('SIGTERM', gracefulShutdown);

initialize().then(() => {
  const server = app.listen(port, hostname, () => {
    console.log(`========================================`);
    console.log(`HTTP & WebSocket 服务器已启动`);
    console.log(`服务地址: http://${hostname}:${port}/`);
    console.log(`========================================`);
    
    createWebSocketServer(server);
  });
});