import { WebSocketServer } from 'ws';
import { Server } from 'http';
import jwt from 'jsonwebtoken';

interface WebSocketClient {
  id: string;
  userId: string;
  role: string;
  ws: any;
}

export class WebSocketService {
  private wss: WebSocketServer;
  private clients: Map<string, WebSocketClient> = new Map();

  constructor(server: Server) {
    this.wss = new WebSocketServer({ 
      server,
      path: '/ws'
    });

    this.wss.on('connection', (ws, req) => {
      this.handleConnection(ws, req);
    });
  }

  private handleConnection(ws: any, req: any) {
    const url = new URL(req.url!, `http://${req.headers.host}`);
    const token = url.searchParams.get('token');

    if (!token) {
      ws.close(1008, 'Token required');
      return;
    }

    try {
      // Simplified token verification for development
      // In development, we'll accept simple user info from query params
      let userId, role;
      
      if (token && token !== 'undefined') {
        try {
          const payload = jwt.decode(token) as any;
          userId = payload?.userId;
          role = payload?.role;
        } catch (e) {
          // Token parsing failed, try query params
        }
      }
      
      // Fallback to query params for development
      if (!userId) {
        userId = url.searchParams.get('userId');
        role = url.searchParams.get('role') || 'client';
      }
      
      if (!userId) {
        console.log('WebSocket: No user ID provided');
        ws.close(1008, 'User ID required');
        return;
      }

      const clientId = Math.random().toString(36).substring(7);
      const client: WebSocketClient = {
        id: clientId,
        userId: userId,
        role: role,
        ws
      };

      this.clients.set(clientId, client);

      console.log(`WebSocket client connected: ${userId} (${role})`);

      // Send welcome message
      this.sendToClient(clientId, {
        type: 'connected',
        message: 'Connected to MediFlash SOS+ real-time updates'
      });

      // Handle incoming messages
      ws.on('message', (data: Buffer) => {
        try {
          const message = JSON.parse(data.toString());
          this.handleMessage(clientId, message);
        } catch (error) {
          console.error('Invalid WebSocket message:', error);
        }
      });

      // Handle disconnection
      ws.on('close', () => {
        this.clients.delete(clientId);
        console.log(`WebSocket client disconnected: ${userId}`);
      });

      ws.on('error', (error: Error) => {
        console.error('WebSocket error:', error);
        this.clients.delete(clientId);
      });

    } catch (error) {
      console.error('WebSocket authentication error:', error);
      ws.close(1008, 'Authentication failed');
    }
  }

  private handleMessage(clientId: string, message: any) {
    const client = this.clients.get(clientId);
    if (!client) return;

    switch (message.type) {
      case 'ping':
        this.sendToClient(clientId, { type: 'pong' });
        break;
      
      case 'subscribe_order':
        // Subscribe to specific order updates
        if (message.orderId) {
          client.ws.orderId = message.orderId;
        }
        break;
      
      case 'location_update':
        // Handle courier location updates
        if (client.role === 'courier' && message.location) {
          this.broadcastCourierLocation(client.userId, message.location);
        }
        break;
      
      default:
        console.warn('Unknown WebSocket message type:', message.type);
    }
  }

  // Send message to specific client
  public sendToClient(clientId: string, message: any) {
    const client = this.clients.get(clientId);
    if (client && client.ws.readyState === 1) {
      client.ws.send(JSON.stringify(message));
    }
  }

  // Send message to specific user
  public sendToUser(userId: string, message: any) {
    this.clients.forEach((client, clientId) => {
      if (client.userId === userId) {
        this.sendToClient(clientId, message);
      }
    });
  }

  // Send message to users by role
  public sendToRole(role: string, message: any) {
    this.clients.forEach((client, clientId) => {
      if (client.role === role) {
        this.sendToClient(clientId, message);
      }
    });
  }

  // Broadcast order update to relevant users
  public broadcastOrderUpdate(orderId: string, orderData: any, userIds: string[]) {
    const message = {
      type: 'order_update',
      orderId,
      data: orderData,
      timestamp: new Date().toISOString()
    };

    userIds.forEach(userId => {
      this.sendToUser(userId, message);
    });
  }

  // Broadcast new order to nearby pharmacies
  public broadcastNewOrderToPharmacies(orderData: any, pharmacyIds: string[]) {
    const message = {
      type: 'new_order',
      data: orderData,
      timestamp: new Date().toISOString()
    };

    pharmacyIds.forEach(pharmacyId => {
      this.sendToUser(pharmacyId, message);
    });
  }

  // Broadcast courier location to order participants
  public broadcastCourierLocation(courierId: string, location: { lat: number; lng: number }) {
    const message = {
      type: 'courier_location',
      courierId,
      location,
      timestamp: new Date().toISOString()
    };

    // Send to all clients subscribed to orders with this courier
    this.clients.forEach((client, clientId) => {
      if (client.ws.orderId) {
        // In a real implementation, check if this courier is assigned to the order
        this.sendToClient(clientId, message);
      }
    });
  }

  // Notify order status change
  public notifyOrderStatusChange(orderId: string, newStatus: string, userIds: string[], details?: any) {
    const message = {
      type: 'order_status_change',
      orderId,
      status: newStatus,
      details,
      timestamp: new Date().toISOString()
    };

    userIds.forEach(userId => {
      this.sendToUser(userId, message);
    });
  }

  // Get connected users count by role
  public getConnectedUsersCount(): { total: number; byRole: Record<string, number> } {
    const byRole: Record<string, number> = {};
    
    this.clients.forEach(client => {
      byRole[client.role] = (byRole[client.role] || 0) + 1;
    });

    return {
      total: this.clients.size,
      byRole
    };
  }

  // Check if user is connected
  public isUserConnected(userId: string): boolean {
    for (const client of this.clients.values()) {
      if (client.userId === userId) {
        return true;
      }
    }
    return false;
  }
}

let wsService: WebSocketService;

export function initWebSocket(server: Server): WebSocketService {
  wsService = new WebSocketService(server);
  return wsService;
}

export function getWebSocketService(): WebSocketService {
  return wsService;
}