server.js

import { createServer } from 'http'
import next from 'next'
import { Server } from 'socket.io'
import { verifyJwtToken } from './lib/auth/jwt.js'
import { wsRateLimiter } from './lib/security/rate-limiter.ts'

const dev = process.env.NODE_ENV !== 'production'
const app = next({ dev })
const handle = app.getRequestHandler()

// WebSocket Service Implementation
class WebSocketService {
  constructor(server) {
    this.io = new Server(server, {
      cors: {
        origin: process.env.NEXTAUTH_URL || 'http://localhost:3000',
        methods: ['GET', 'POST'],
        credentials: true
      },
      transports: ['websocket', 'polling']
    })

    this.userSockets = new Map()
    this.conversationRooms = new Map()
    this.setupEventHandlers()
  }

  async authenticateSocket(socket) {
    try {
      const token = socket.handshake.auth.token ||
                    (socket.handshake.headers.authorization || socket.request.headers.authorization || '').replace('Bearer ', '')
      
      if (!token) {
        console.warn('WebSocket authentication failed: No token provided')
        throw new Error('No authentication token provided')
      }

      // Validate token format
      if (typeof token !== 'string' || token.length < 10) {
        console.warn('WebSocket authentication failed: Invalid token format')
        throw new Error('Invalid token format')
      }

      // Verify the JWT token generated by our auth endpoint
      const session = await verifyJwtToken(token, process.env.AUTH_SECRET)

      if (!session || !session.sub) {
        console.warn('WebSocket authentication failed: Invalid session')
        throw new Error('Invalid authentication token')
      }

      // Check if this is an anonymous user
      const isAnonymous = session.sub && session.sub.startsWith('anon-')
      
      console.log(`✅ WebSocket authentication successful for ${isAnonymous ? 'anonymous' : 'authenticated'} user: ${session.email || session.sub}`)
      
      // Store user info for later use
      socket.userId = session.sub
      socket.userEmail = session.email
      socket.userRole = session.role || (isAnonymous ? 'anonymous' : 'subscriber')
      socket.isAnonymous = isAnonymous
      
      return session
    } catch (error) {
      console.error('❌ Socket authentication failed:', error.message)
      // SECURITY: Don't expose detailed error messages to clients
      throw new Error('Authentication failed')
    }
  }

  setupEventHandlers() {
    // SECURITY: Rate limiting middleware - prevent connection flooding
    this.io.use((socket, next) => {
      const identifier = socket.handshake.address || 
                        socket.request.connection?.remoteAddress || 
                        'unknown'
      
      if (wsRateLimiter.isRateLimited(identifier)) {
        console.warn(`⚠️ WebSocket rate limit exceeded for ${identifier}`)
        const error = new Error('Too many connection attempts')
        error.data = { code: 'RATE_LIMIT_EXCEEDED', retryAfter: 60 }
        return next(error)
      }
      
      next()
    })
    
    // Authentication middleware
    this.io.use(async (socket, next) => {
      try {
        const session = await this.authenticateSocket(socket)
        socket.userId = session.sub
        socket.userEmail = session.email || session.sub
        socket.sessionData = session
        next()
      } catch (error) {
        console.error('WebSocket middleware authentication failed:', error.message)
        // SECURITY: Generic error message to prevent information leakage
        next(new Error('Connection failed'))
      }
    })

    this.io.on('connection', (socket) => {
      try {
        console.log(`User connected: ${socket.userEmail} (${socket.userId})`)
        
        // Store user socket mapping
        this.userSockets.set(socket.userId, socket)

        // Join user to their personal room for direct notifications
        socket.join(`user:${socket.userId}`)

        // Handle conversation joining with SECURITY validation
        socket.on('join_conversation', (conversationId) => {
          try {
            // SECURITY: Input validation and sanitization
            if (!conversationId || typeof conversationId !== 'string') {
              socket.emit('error', { message: 'Invalid conversation ID' })
              return
            }
            
            // SECURITY: Validate conversation ID format (alphanumeric + hyphens, max 100 chars)
            const sanitizedId = conversationId.trim().substring(0, 100)
            if (!/^[a-zA-Z0-9_-]+$/.test(sanitizedId)) {
              console.warn(`⚠️ Invalid conversation ID format attempted by ${socket.userEmail}: ${conversationId}`)
              socket.emit('error', { message: 'Invalid conversation ID format' })
              return
            }

            socket.join(`conversation:${sanitizedId}`)
            
            // Track users in conversation room
            if (!this.conversationRooms.has(sanitizedId)) {
              this.conversationRooms.set(sanitizedId, new Set())
            }
            this.conversationRooms.get(sanitizedId).add(socket.userId)
            
            console.log(`User ${socket.userEmail} joined conversation: ${sanitizedId}`)
            socket.emit('conversation_joined', { conversationId: sanitizedId })
          } catch (error) {
            console.error('Error joining conversation:', error)
            socket.emit('error', { message: 'Failed to join conversation' })
          }
        })

        // Handle conversation leaving with error handling
        socket.on('leave_conversation', (conversationId) => {
          try {
            if (!conversationId || typeof conversationId !== 'string') {
              socket.emit('error', { message: 'Invalid conversation ID' })
              return
            }

            socket.leave(`conversation:${conversationId}`)
            
            if (this.conversationRooms.has(conversationId)) {
              this.conversationRooms.get(conversationId).delete(socket.userId)
              if (this.conversationRooms.get(conversationId).size === 0) {
                this.conversationRooms.delete(conversationId)
              }
            }
            
            console.log(`User ${socket.userEmail} left conversation: ${conversationId}`)
            socket.emit('conversation_left', { conversationId })
          } catch (error) {
            console.error('Error leaving conversation:', error)
            socket.emit('error', { message: 'Failed to leave conversation' })
          }
        })

        // Heartbeat/Keepalive handling for modern WebSocket connection
        socket.on('ping', () => {
          socket.emit('pong')
        })

        // Notification subscription management
        socket.on('subscribe', ({ topic }) => {
          socket.join(topic)
          console.log(`Socket ${socket.id} subscribed to ${topic}`)
          
          // Join user-specific notification channel
          if (topic === 'user:notifications' && socket.userId) {
            socket.join(`user:${socket.userId}:notifications`)
          }
        })

        socket.on('unsubscribe', ({ topic }) => {
          socket.leave(topic)
          console.log(`Socket ${socket.id} unsubscribed from ${topic}`)
        })

        // Notification operations - Replace polling with push
        socket.on('notification:get_count', async () => {
          try {
            // This would connect to your Firebase/database to get unread count
            // For now, sending mock data
            const unreadCount = Math.floor(Math.random() * 10)
            socket.emit('notification:unread_count', unreadCount)
          } catch (error) {
            console.error('Error fetching notification count:', error)
          }
        })

        socket.on('notification:mark_read', async ({ ids }) => {
          try {
            // Mark notifications as read in database
            console.log(`Marking notifications as read: ${ids.join(', ')}`)
            // Update database here
            
            // Emit updated count
            socket.emit('notification:unread_count', 0)
          } catch (error) {
            console.error('Error marking notifications as read:', error)
          }
        })

        // Handle message sending with SECURITY validation
        socket.on('send_message', async (messageData) => {
          try {
            const { conversationId, content, type = 'text' } = messageData
            
            // SECURITY: Validate required fields
            if (!conversationId || !content) {
              socket.emit('error', { message: 'Missing required message data' })
              return
            }
            
            // SECURITY: Validate and sanitize conversation ID
            const sanitizedConvId = conversationId.toString().trim().substring(0, 100)
            if (!/^[a-zA-Z0-9_-]+$/.test(sanitizedConvId)) {
              socket.emit('error', { message: 'Invalid conversation ID' })
              return
            }
            
            // SECURITY: Validate message type
            const allowedTypes = ['text', 'image', 'file', 'system']
            const sanitizedType = allowedTypes.includes(type) ? type : 'text'
            
            // SECURITY: Sanitize and validate content (prevent XSS)
            const sanitizedContent = content.toString().substring(0, 10000) // Max 10K chars
              .replace(/[<>]/g, '') // Basic XSS prevention (use DOMPurify in production)
            
            if (sanitizedContent.length === 0) {
              socket.emit('error', { message: 'Invalid message content' })
              return
            }

            // Create message object with server-side data
            const message = {
              id: `msg_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`,
              conversationId: sanitizedConvId,
              senderId: socket.userId,
              senderName: socket.userEmail,
              content: sanitizedContent,
              type: sanitizedType,
              timestamp: new Date(),
              status: 'sent'
            }

            // Broadcast to all users in the conversation
            this.io.to(`conversation:${sanitizedConvId}`).emit('message_received', message)
            
            console.log(`Message sent in conversation ${sanitizedConvId} by ${socket.userEmail}`)
            
          } catch (error) {
            console.error('Error handling message:', error)
            socket.emit('error', { message: 'Failed to send message' })
          }
        })

        // Handle typing indicators
        socket.on('typing_start', (conversationId) => {
          socket.to(`conversation:${conversationId}`).emit('user_typing', {
            userId: socket.userId,
            userEmail: socket.userEmail,
            isTyping: true
          })
        })

        socket.on('typing_stop', (conversationId) => {
          socket.to(`conversation:${conversationId}`).emit('user_typing', {
            userId: socket.userId,
            userEmail: socket.userEmail,
            isTyping: false
          })
        })

        // Handle presence updates
        socket.on('update_presence', (status) => {
          socket.broadcast.emit('presence_update', {
            userId: socket.userId,
            userEmail: socket.userEmail,
            isOnline: status === 'online',
            lastSeen: new Date()
          })
        })

        // NEW: Handle token refresh without disconnection - PREVENTS SUBSCRIPTION LOOPS
        socket.on('auth:refresh', async ({ token }) => {
          try {
            console.log(`🔄 Token refresh request from ${socket.userEmail}`)
            
            // SECURITY: Validate token format before processing
            if (!token || typeof token !== 'string' || token.length < 10) {
              throw new Error('Invalid token format')
            }
            
            // Validate new token using the same verification logic
            const session = await verifyJwtToken(token, process.env.AUTH_SECRET)
            
            // SECURITY FIX: Validate token expiry
            const now = Math.floor(Date.now() / 1000)
            if (!session.exp || session.exp < now) {
              console.warn(`⚠️ Token refresh attempted with expired token for ${socket.userEmail}`)
              throw new Error('Token has expired')
            }
            
            // SECURITY FIX: Validate token not issued in the future (clock skew protection)
            if (session.iat && session.iat > now + 60) { // Allow 60 seconds clock skew
              console.warn(`⚠️ Token refresh attempted with future-dated token for ${socket.userEmail}`)
              throw new Error('Invalid token timestamp')
            }
            
            if (session && session.sub === socket.userId) {
              // Update socket auth data without disconnection
              socket.sessionData = session
              socket.userEmail = session.email || socket.userEmail
              socket.userRole = session.role || socket.userRole
              
              // Confirm successful refresh to client
              socket.emit('auth:refreshed', { 
                success: true,
                message: 'Token refreshed successfully',
                expiresIn: session.exp - now
              })
              
              console.log(`✅ Token refreshed successfully for ${socket.userEmail}`)
            } else {
              // Invalid token or user mismatch
              console.warn(`❌ Token refresh failed for ${socket.userEmail}: Invalid token or user mismatch`)
              socket.emit('auth:refreshed', { 
                success: false, 
                error: 'Invalid token or user mismatch'
              })
              
              // Force disconnection for security
              setTimeout(() => {
                socket.disconnect(true)
              }, 1000)
            }
          } catch (error) {
            console.error(`❌ Token refresh error for ${socket.userEmail}:`, error.message)
            socket.emit('auth:refreshed', { 
              success: false, 
              error: 'Token validation failed'
            })
            
            // Force disconnection on validation failure
            setTimeout(() => {
              socket.disconnect(true)
            }, 1000)
          }
        })

        // Handle disconnection
        socket.on('disconnect', (reason) => {
          console.log(`User disconnected: ${socket.userEmail} (${reason})`)
          
          // Remove from user sockets mapping
          this.userSockets.delete(socket.userId)
          
          // Remove from all conversation rooms
          for (const [conversationId, users] of this.conversationRooms.entries()) {
            if (users.has(socket.userId)) {
              users.delete(socket.userId)
              if (users.size === 0) {
                this.conversationRooms.delete(conversationId)
              }
            }
          }

          // Broadcast offline status
          socket.broadcast.emit('presence_update', {
            userId: socket.userId,
            userEmail: socket.userEmail,
            isOnline: false,
            lastSeen: new Date()
          })
        })

        // Send connection confirmation
        socket.emit('connected', { 
          message: 'WebSocket connected successfully',
          userId: socket.userId,
          userEmail: socket.userEmail
        })
      } catch (error) {
        console.error('Error in connection handler:', error)
        socket.emit('error', { message: 'Connection handler error' })
      }
    })
  }

  // Public API methods
  sendToUser(userId, event, data) {
    const socket = this.userSockets.get(userId)
    if (socket) {
      socket.emit(event, data)
      return true
    }
    return false
  }

  sendToConversation(conversationId, event, data) {
    this.io.to(`conversation:${conversationId}`).emit(event, data)
  }

  broadcast(event, data) {
    this.io.emit(event, data)
  }

  getActiveUsers() {
    return Array.from(this.userSockets.keys())
  }

  getConversationUsers(conversationId) {
    return Array.from(this.conversationRooms.get(conversationId) || [])
  }
}

// Start the server
app.prepare().then(() => {
  const server = createServer((req, res) => {
    handle(req, res)
  })

  // Initialize WebSocket service
  const wsService = new WebSocketService(server)
  
  // Make wsService available globally for API routes if needed
  global.wsService = wsService

  const port = process.env.PORT || 3000
  
  server.listen(port, (err) => {
    if (err) throw err
    console.log(`🚀 Server ready on http://localhost:${port}`)
    console.log(`📡 WebSocket server ready`)
  })
}).catch((ex) => {
  console.error('Error starting server:', ex)
  process.exit(1)
})