import WebSocket from 'ws';
import http from 'http';
import { createApp } from '../src/app';
import { WebSocketHandler } from '../src/websocket/wsHandler';
import { TaskBrokerWebSocketServer } from '../src/websocket/wsServer';

async function runTestSimulation() {
  console.log('--- Starting Distributed Task Broker Simulation Test ---');

  // Setup app & server on test port 7899
  const context = createApp();
  const server = http.createServer(context.app);
  const wsHandler = new WebSocketHandler(
    context.workerService,
    context.taskQueueService,
    context.progressService,
    context.resultService
  );
  const wsServer = new TaskBrokerWebSocketServer(server, wsHandler);
  context.adminController.setWsServer(wsServer);

  const PORT = 7899;
  await new Promise<void>((resolve) => server.listen(PORT, resolve));
  console.log(`Test server running on port ${PORT}`);

  // Create initial batch tasks (100,000 rows, 10,000 chunkSize = 10 tasks)
  await context.taskQueueService.generateBatchTasks(100000, 10000);

  // Helper to create a simulated worker
  const simulateWorker = (id: string, name: string, platform: 'colab' | 'kaggle' | 'local', gpuName: string) => {
    return new Promise<void>((resolve, reject) => {
      const ws = new WebSocket(`ws://localhost:${PORT}/ws`);

      ws.on('open', () => {
        console.log(`[${id}] WebSocket connected. Registering worker...`);
        ws.send(
          JSON.stringify({
            type: 'register',
            payload: {
              workerId: id,
              workerName: name,
              platform: platform,
              specs: { cpu: 8, ram: '32GB', gpu: gpuName, batchSize: 64 },
            },
          })
        );
      });

      ws.on('message', (data: WebSocket.RawData) => {
        const msg = JSON.parse(data.toString());
        // console.log(`[${id}] Received server message:`, msg.type);

        if (msg.type === 'registered') {
          console.log(`[${id}] Registered successfully! Requesting first OCR task...`);
          ws.send(JSON.stringify({ type: 'request_task', payload: { workerId: id } }));
        } else if (msg.type === 'task_assignment') {
          const task = msg.payload;
          console.log(
            `[${id}] ASSIGNED TASK: ${task.taskId} | Rows: ${task.startRow} -> ${task.endRow} (${task.totalRows} rows) | Offsets: ${task.startOffset} -> ${task.endOffset}`
          );

          // Send heartbeat and progress update
          ws.send(
            JSON.stringify({
              type: 'progress',
              payload: { workerId: id, taskId: task.taskId, progress: 50, rowsProcessed: task.totalRows / 2, currentSpeed: 250 },
            })
          );

          // Complete task after short delay
          setTimeout(() => {
            console.log(`[${id}] Completing task ${task.taskId}...`);
            ws.send(
              JSON.stringify({
                type: 'task_completed',
                payload: { workerId: id, taskId: task.taskId, rowsProcessed: task.totalRows },
              })
            );
          }, 300);
        } else if (msg.type === 'wait') {
          console.log(`[${id}] Queue empty or wait state received: ${msg.payload.reason}`);
          ws.close();
          resolve();
        }
      });

      ws.on('error', (err) => {
        console.error(`[${id}] WS Error:`, err);
        reject(err);
      });
    });
  };

  // Run 3 concurrent remote workers
  await Promise.all([
    simulateWorker('colab-worker-1', 'Google Colab T4 Node', 'colab', 'NVIDIA T4'),
    simulateWorker('kaggle-worker-1', 'Kaggle P100 Node', 'kaggle', 'NVIDIA P100'),
    simulateWorker('local-worker-1', 'Desktop Workstation', 'local', 'RTX 4090'),
  ]);

  // Check system statistics after simulation
  const stats = await context.statisticsService.getSystemStatistics();
  console.log('\n--- Final System Statistics ---');
  console.log(`Total Tasks: ${stats.tasks.total}`);
  console.log(`Completed Tasks: ${stats.tasks.completed}`);
  console.log(`Pending Tasks: ${stats.tasks.pending}`);
  console.log(`Total Rows Processed: ${stats.throughput.totalRowsProcessed}`);
  console.log(`Completion Percentage: ${stats.tasks.completionPercentage}%`);

  server.close();
  console.log('--- Simulation Completed Successfully ---');
}

runTestSimulation().catch(console.error);