import { Server as HttpServer } from 'http';
import { WebSocketServer, WebSocket } from 'ws';
import jwt from 'jsonwebtoken';
import { config } from '../config/index.js';
import { logger } from '../utils/logger.js';
import type { JwtPayload } from '../middleware/auth.js';
import { WS_EVENTS } from '@media-manager/shared';

interface AuthenticatedWebSocket extends WebSocket {
  userId?: number;
  isAlive?: boolean;
}

class WebSocketManager {
  private wss: WebSocketServer | null = null;
  private clients: Map<number, Set<AuthenticatedWebSocket>> = new Map();
  // 验证码回调监听器
  private captchaListeners: Map<string, (code: string) => void> = new Map();

  setup(server: HttpServer): void {
    this.wss = new WebSocketServer({ server, path: '/ws' });

    this.wss.on('connection', (ws: AuthenticatedWebSocket, req) => {
      logger.info(`[WS] New connection from ${req.socket.remoteAddress}, URL: ${req.url}`);
      ws.isAlive = true;

      ws.on('pong', () => {
        ws.isAlive = true;
      });

      ws.on('message', (data) => {
        logger.info(`[WS] Received message: ${data.toString().slice(0, 200)}`);
        this.handleMessage(ws, data.toString());
      });

      ws.on('close', (code, reason) => {
        logger.info(`[WS] Connection closed, code: ${code}, reason: ${reason?.toString()}`);
        this.removeClient(ws);
      });

      ws.on('error', (error) => {
        logger.error('[WS] Connection error:', error);
        this.removeClient(ws);
      });

      // 发送连接成功消息，等待认证
      logger.info('[WS] Sending connection message...');
      this.send(ws, { type: 'connection', payload: { message: 'Connected, please authenticate' } });
    });

    // 心跳检测
    setInterval(() => {
      this.wss?.clients.forEach((ws: AuthenticatedWebSocket) => {
        if (!ws.isAlive) {
          this.removeClient(ws);
          return ws.terminate();
        }
        ws.isAlive = false;
        ws.ping();
      });
    }, 30000);

    logger.info('WebSocket server initialized');
  }

  private handleMessage(ws: AuthenticatedWebSocket, data: string): void {
    try {
      const message = JSON.parse(data);

      switch (message.type) {
        case WS_EVENTS.AUTH:
          this.handleAuth(ws, message.payload?.token);
          break;
        case WS_EVENTS.PING:
          this.send(ws, { type: WS_EVENTS.PONG, timestamp: Date.now() });
          break;
        case WS_EVENTS.CAPTCHA_SUBMIT:
          this.handleCaptchaSubmit(message.payload);
          break;
        default:
          logger.warn('Unknown WebSocket message type:', message.type);
      }
    } catch (error) {
      logger.error('Failed to parse WebSocket message:', error);
    }
  }
  
  /**
   * 处理验证码提交
   */
  private handleCaptchaSubmit(payload: { captchaTaskId: string; code: string }): void {
    const { captchaTaskId, code } = payload || {};
    if (!captchaTaskId || !code) {
      logger.warn('[WS] Invalid captcha submit payload');
      return;
    }
    
    const listener = this.captchaListeners.get(captchaTaskId);
    if (listener) {
      logger.info(`[WS] Captcha submitted for task ${captchaTaskId}`);
      listener(code);
    } else {
      logger.warn(`[WS] No listener for captcha task ${captchaTaskId}`);
    }
  }

  private handleAuth(ws: AuthenticatedWebSocket, token?: string): void {
    logger.info(`[WS] Authenticating, token: ${token ? token.slice(0, 20) + '...' : 'none'}`);
    
    if (!token) {
      logger.warn('[WS] Auth failed: no token');
      this.send(ws, { type: WS_EVENTS.AUTH_ERROR, payload: { message: 'Token required' } });
      return;
    }

    try {
      const decoded = jwt.verify(token, config.jwt.secret) as JwtPayload;
      ws.userId = decoded.userId;

      // 添加到用户的连接列表
      if (!this.clients.has(decoded.userId)) {
        this.clients.set(decoded.userId, new Set());
      }
      this.clients.get(decoded.userId)!.add(ws);

      this.send(ws, { 
        type: WS_EVENTS.AUTH_SUCCESS, 
        payload: { userId: decoded.userId, message: 'Authenticated' } 
      });

      logger.info(`[WS] Auth success for user ${decoded.userId}, total connections: ${this.clients.get(decoded.userId)!.size}`);
    } catch (error) {
      logger.warn('[WS] Auth failed: invalid token', error);
      this.send(ws, { type: WS_EVENTS.AUTH_ERROR, payload: { message: 'Invalid token' } });
    }
  }

  private removeClient(ws: AuthenticatedWebSocket): void {
    if (ws.userId) {
      const userClients = this.clients.get(ws.userId);
      if (userClients) {
        userClients.delete(ws);
        if (userClients.size === 0) {
          this.clients.delete(ws.userId);
        }
      }
    }
  }

  private send(ws: WebSocket, data: object): void {
    if (ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify({ ...data, timestamp: Date.now() }));
    }
  }

  /**
   * 向指定用户的所有客户端发送消息
   */
  sendToUser(userId: number, type: string, payload?: unknown): void {
    const userClients = this.clients.get(userId);
    if (userClients) {
      const message = JSON.stringify({ type, payload, timestamp: Date.now() });
      logger.info(`[WS] Sending to user ${userId}: type=${type}, payload=${JSON.stringify(payload)}`);
      logger.info(`[WS] Full message: ${message}`);
      userClients.forEach((ws) => {
        if (ws.readyState === WebSocket.OPEN) {
          ws.send(message);
        }
      });
    } else {
      logger.warn(`[WS] No clients for user ${userId}, message not sent: type=${type}`);
    }
  }

  /**
   * 向所有已认证的客户端广播消息
   */
  broadcast(type: string, payload?: unknown): void {
    const message = JSON.stringify({ type, payload, timestamp: Date.now() });
    this.clients.forEach((userClients) => {
      userClients.forEach((ws) => {
        if (ws.readyState === WebSocket.OPEN) {
          ws.send(message);
        }
      });
    });
  }

  /**
   * 获取在线用户数
   */
  getOnlineUserCount(): number {
    return this.clients.size;
  }

  /**
   * 检查用户是否在线
   */
  isUserOnline(userId: number): boolean {
    const userClients = this.clients.get(userId);
    return userClients ? userClients.size > 0 : false;
  }
  
  /**
   * 注册验证码提交监听
   */
  onCaptchaSubmit(captchaTaskId: string, callback: (code: string) => void): void {
    this.captchaListeners.set(captchaTaskId, callback);
    logger.info(`[WS] Registered captcha listener for ${captchaTaskId}`);
  }
  
  /**
   * 移除验证码监听
   */
  removeCaptchaListener(captchaTaskId: string): void {
    this.captchaListeners.delete(captchaTaskId);
    logger.info(`[WS] Removed captcha listener for ${captchaTaskId}`);
  }
}

export const wsManager = new WebSocketManager();

export function setupWebSocket(server: HttpServer): void {
  wsManager.setup(server);
}
