import ws from 'k6/ws';
import { check, sleep } from 'k6';
import http from 'k6/http';
import encoding from 'k6/encoding';

// Test different connection scales
const SCALE = __ENV.SCALE || '1k'; // 1k, 10k, 100k

const scales = {
  '1k': { vus: 1000, duration: '5m' },
  '10k': { vus: 10000, duration: '10m' },
  '100k': { vus: 100000, duration: '15m' },
};

export const options = scales[SCALE] || scales['1k'];

const BASE_URL = __ENV.BASE_URL || 'http://gotify-256:80';
const USERNAME = __ENV.USERNAME || 'admin';
const PASSWORD = __ENV.PASSWORD || 'admin';

let authToken = null;
let clientToken = null;

export function setup() {
  const baseUrl = BASE_URL.replace(/^http/, 'http');
  
  const credentials = encoding.b64encode(`${USERNAME}:${PASSWORD}`);
  
  const clientRes = http.post(`${baseUrl}/client`, JSON.stringify({
    name: `k6-scale-${__VU}-${Date.now()}`,
  }), {
    headers: {
      'Content-Type': 'application/json',
      'Authorization': `Basic ${credentials}`,
    },
  });
  
  if (clientRes.status !== 200) {
    console.error('Failed to create client:', clientRes.status, clientRes.body);
    return {};
  }
  
  clientToken = JSON.parse(clientRes.body).token;
  
  return { clientToken, baseUrl };
}

export default function (data) {
  if (!data.clientToken) return;

  const wsUrl = data.baseUrl.replace(/^http/, 'ws') + '/stream?token=' + data.clientToken;
  
  const startTime = Date.now();
  const response = ws.connect(wsUrl, {}, function (socket) {
    const connectTime = Date.now() - startTime;
    
    socket.on('open', () => {
      // Track connection time
      console.log(`VU ${__VU}: Connected in ${connectTime}ms`);
    });
    
    socket.on('message', (data) => {
      const msg = JSON.parse(data);
      const receiveTime = Date.now() - startTime;
      check(msg, {
        'message received': (m) => m.id !== undefined,
        'receive latency acceptable': () => receiveTime < 5000,
      });
    });
    
    socket.on('close', () => {
      console.log(`VU ${__VU}: Connection closed after ${Date.now() - startTime}ms`);
    });
    
    socket.on('error', (e) => {
      if (e.error() !== 'websocket: close sent') {
        console.error(`VU ${__VU}: Error:`, e.error());
      }
    });
    
    // Keep connection alive
    sleep(60);
  });

  check(response, {
    'connection successful': (r) => r && r.status === 101,
    'connection time < 1s': () => (Date.now() - startTime) < 1000,
  });

  sleep(1);
}