import { Elysia, t } from "elysia";
import { StateManager } from '../services/StateManager';
import { BedService } from '../services/BedService';
import { WebSocketMessage, StateUpdateEvent } from '../types/FrontendState';

// Define WebSocket client type
interface WSClient {
  id: string;
  send: (message: string) => void;
  readyState: number;
}

// Singleton instances
let stateManager: StateManager | null = null;
let bedService: BedService | null = null;
const clients = new Map<string, WSClient>();

// Initialize services
async function initializeServices() {
  if (!bedService) {
    bedService = new BedService();
    await bedService.initialize();
  }
  
  if (!stateManager) {
    stateManager = new StateManager(bedService);
    await stateManager.initializeState();
    
    // Listen for state updates and broadcast to clients
    stateManager.on('stateUpdate', (event: StateUpdateEvent) => {
      broadcastStateUpdate(event);
    });
  }
}

// Broadcast state updates to all connected clients
function broadcastStateUpdate(event: StateUpdateEvent) {
  const message: WebSocketMessage = {
    type: 'STATE_UPDATE',
    payload: event,
    timestamp: new Date()
  };

  const messageStr = JSON.stringify(message);
  const deadClients: string[] = [];

  clients.forEach((ws, id) => {
    try {
      if (ws.readyState === 1) { // WebSocket.OPEN
        ws.send(messageStr);
      } else {
        deadClients.push(id);
      }
    } catch (error) {
      console.error('Failed to send message to client:', error);
      deadClients.push(id);
    }
  });

  // Clean up dead clients
  deadClients.forEach(id => {
    clients.delete(id);
  });

  // Update connection count
  if (stateManager) {
    stateManager.updateConnectionCount(clients.size);
  }
}

// Send current state to a specific client
async function sendCurrentState(ws: WSClient) {
  try {
    await initializeServices();
    
    if (stateManager) {
      const state = stateManager.getState();
      const message: WebSocketMessage = {
        type: 'STATE_UPDATE',
        payload: {
          type: 'FULL_STATE',
          timestamp: new Date(),
          data: state
        },
        timestamp: new Date()
      };

      ws.send(JSON.stringify(message));
    }
  } catch (error) {
    console.error('Failed to send current state:', error);
  }
}

const stateRouter = new Elysia()
  .ws('/ws', {
    body: t.Object({
      type: t.Union([
        t.Literal('ACKNOWLEDGE_ALERT'),
        t.Literal('SILENCE_ALERT'),
        t.Literal('HEARTBEAT')
      ]),
      payload: t.Optional(t.Object({
        alertId: t.Optional(t.String())
      }))
    }),
    
    open: async (ws) => {
      console.log('WebSocket client connected:', ws.id);
      clients.set(ws.id, ws);
      
      // Send current state to new client
      await sendCurrentState(ws);
      
      // Update connection count
      await initializeServices();
      if (stateManager) {
        stateManager.updateConnectionCount(clients.size);
      }
    },
    
    message: async (ws, message) => {
      try {
        await initializeServices();
        
        switch (message.type) {
          case 'ACKNOWLEDGE_ALERT':
            if (message.payload?.alertId && stateManager) {
              await stateManager.acknowledgeAlert(message.payload.alertId);
            }
            break;
            
          case 'SILENCE_ALERT':
            if (message.payload?.alertId && stateManager) {
              await stateManager.silenceAlert(message.payload.alertId);
            }
            break;
            
          case 'HEARTBEAT':
            // Respond to heartbeat
            ws.send(JSON.stringify({
              type: 'HEARTBEAT',
              payload: { message: 'pong' },
              timestamp: new Date()
            }));
            break;
            
          default:
            console.warn('Unknown WebSocket message type:', message.type);
        }
      } catch (error) {
        console.error('Error handling client message:', error);
        ws.send(JSON.stringify({
          type: 'ERROR',
          payload: { message: 'Invalid message format' },
          timestamp: new Date()
        }));
      }
    },
    
    close: (ws) => {
      console.log('WebSocket client disconnected:', ws.id);
      clients.delete(ws.id);
      
      // Update connection count
      if (stateManager) {
        stateManager.updateConnectionCount(clients.size);
      }
    }
  });

export default stateRouter;