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

// Test configuration
export const options = {
  stages: [
    { duration: '30s', target: 1000 },   // Ramp up to 1000 connections
    { duration: '1m', target: 1000 },     // Stay at 1000 connections
    { duration: '30s', target: 5000 },   // Ramp up to 5000 connections
    { duration: '2m', target: 5000 },    // Stay at 5000 connections
    { duration: '30s', target: 10000 },  // Ramp up to 10000 connections
    { duration: '2m', target: 10000 },   // Stay at 10000 connections
    { duration: '30s', target: 0 },      // Ramp down
  ],
  thresholds: {
    'ws_connecting': ['p(95)<500'],      // 95% of connections should connect in <500ms
    'ws_session_duration': ['p(95)<1000'], // 95% of sessions should last <1s
    'ws_ping': ['p(95)<100'],            // 95% of pings should be <100ms
    'ws_messages_sent': ['rate>100'],     // Should send >100 messages/sec
    'ws_messages_received': ['rate>100'], // Should receive >100 messages/sec
  },
};

// Configuration - can be overridden via environment variables
const BASE_URL = __ENV.BASE_URL || 'http://gotify-256:80';
const USERNAME = __ENV.USERNAME || 'admin';
const PASSWORD = __ENV.PASSWORD || 'admin';
const CONNECTIONS_PER_VU = __ENV.CONNECTIONS_PER_VU || 1;

// Global variables for authentication
let clientToken = null;

// Setup: Create a client token using basic auth
export function setup() {
  const baseUrl = BASE_URL.replace(/^http/, 'http');
  
  // Encode basic auth credentials
  const credentials = encoding.b64encode(`${USERNAME}:${PASSWORD}`);
  
  // Create a client token for WebSocket connection using basic auth
  const clientRes = http.post(`${baseUrl}/client`, JSON.stringify({
    name: `k6-client-${__VU}-${__ITER}-${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: clientToken,
    baseUrl: baseUrl,
  };
}

export default function (data) {
  if (!data.clientToken) {
    console.error('No client token available');
    return;
  }

  const wsUrl = data.baseUrl.replace(/^http/, 'ws') + '/stream?token=' + data.clientToken;
  const params = { tags: { name: 'WebSocket Stream' } };

  const response = ws.connect(wsUrl, params, function (socket) {
    socket.on('open', function () {
      // Connection established
      console.log(`VU ${__VU}: WebSocket connection opened`);
    });

    socket.on('message', function (data) {
      // Message received
      const message = JSON.parse(data);
      check(message, {
        'message received': (msg) => msg.id !== undefined,
        'message has content': (msg) => msg.message !== undefined,
      });
    });

    socket.on('ping', function () {
      // Server ping received
    });

    socket.on('pong', function () {
      // Pong received
    });

    socket.on('close', function () {
      // Connection closed
      console.log(`VU ${__VU}: WebSocket connection closed`);
    });

    socket.on('error', function (e) {
      if (e.error() !== 'websocket: close sent') {
        console.error(`VU ${__VU}: WebSocket error:`, e.error());
      }
    });

    // Keep connection alive for the duration of the test
    // The server will send ping messages periodically
    sleep(60); // Keep connection for 60 seconds
  });

  check(response, {
    'status is 101': (r) => r && r.status === 101,
    'protocol is websocket': (r) => r && r.url && r.url.startsWith('ws'),
  });

  sleep(1);
}