import * as plugins from '../ts_server/plugins.js';
import { expect, tap } from '@git.zone/tstest/tapbundle';
import { tapNodeTools } from '@git.zone/tstest/tapbundle_node';
import * as dnsPacket from 'dns-packet';
import * as dgram from 'dgram';

import * as smartdns from '../ts_server/index.js';

let dnsServer: smartdns.DnsServer;

// Port management for tests
let nextHttpsPort = 8200;
let nextUdpPort = 8201;

function getUniqueHttpsPort() {
  return nextHttpsPort++;
}

function getUniqueUdpPort() {
  return nextUdpPort++;
}

// Cleanup function for servers
async function stopServer(server: smartdns.DnsServer | null | undefined) {
  if (!server) {
    return;
  }
  
  try {
    const stopPromise = server.stop();
    const timeoutPromise = new Promise((_, reject) => {
      setTimeout(() => reject(new Error('Stop operation timed out')), 5000);
    });
    
    await Promise.race([stopPromise, timeoutPromise]);
  } catch (e) {
    console.log('Handled error when stopping server:', e.message || e);
    
    // Force close if normal stop fails
    try {
      // @ts-ignore - accessing private properties for emergency cleanup
      if (server.httpsServer) {
        (server as any).httpsServer.close();
        (server as any).httpsServer = null;
      }
      // @ts-ignore - accessing private properties for emergency cleanup
      if (server.udpServer) {
        (server as any).udpServer.close();
        (server as any).udpServer = null;
      }
    } catch (forceError) {
      console.log('Force cleanup error:', forceError.message || forceError);
    }
  }
}

tap.test('should demonstrate the current limitation with multiple NS records', async () => {
  const httpsData = await tapNodeTools.createHttpsCert();
  const udpPort = getUniqueUdpPort();
  
  dnsServer = new smartdns.DnsServer({
    httpsKey: httpsData.key,
    httpsCert: httpsData.cert,
    httpsPort: getUniqueHttpsPort(),
    udpPort: udpPort,
    dnssecZone: 'example.com',
  });
  
  // Register multiple NS record handlers for the same domain
  dnsServer.registerHandler('example.com', ['NS'], (question) => {
    console.log('First NS handler called');
    return {
      name: question.name,
      type: 'NS',
      class: 'IN',
      ttl: 3600,
      data: 'ns1.example.com',
    };
  });
  
  dnsServer.registerHandler('example.com', ['NS'], (question) => {
    console.log('Second NS handler called');
    return {
      name: question.name,
      type: 'NS',
      class: 'IN',
      ttl: 3600,
      data: 'ns2.example.com',
    };
  });
  
  await dnsServer.start();
  
  const client = dgram.createSocket('udp4');
  
  const query = dnsPacket.encode({
    type: 'query',
    id: 1,
    flags: dnsPacket.RECURSION_DESIRED,
    questions: [
      {
        name: 'example.com',
        type: 'NS',
        class: 'IN',
      },
    ],
  });
  
  const responsePromise = new Promise<dnsPacket.Packet>((resolve, reject) => {
    client.on('message', (msg) => {
      const dnsResponse = dnsPacket.decode(msg);
      resolve(dnsResponse);
      client.close();
    });
    
    client.on('error', (err) => {
      reject(err);
      client.close();
    });
    
    client.send(query, udpPort, 'localhost', (err) => {
      if (err) {
        reject(err);
        client.close();
      }
    });
  });
  
  const dnsResponse = await responsePromise;
  
  console.log('Current behavior - NS records returned:', dnsResponse.answers.length);
  console.log('NS records:', dnsResponse.answers.map(a => (a as any).data));
  
  // CURRENT BEHAVIOR: Only returns 1 NS record due to the break statement
  expect(dnsResponse.answers.length).toEqual(1);
  expect((dnsResponse.answers[0] as any).data).toEqual('ns1.example.com');
  
  await stopServer(dnsServer);
  dnsServer = null;
});

tap.test('should demonstrate the limitation with multiple A records (round-robin)', async () => {
  const httpsData = await tapNodeTools.createHttpsCert();
  const udpPort = getUniqueUdpPort();
  
  dnsServer = new smartdns.DnsServer({
    httpsKey: httpsData.key,
    httpsCert: httpsData.cert,
    httpsPort: getUniqueHttpsPort(),
    udpPort: udpPort,
    dnssecZone: 'example.com',
  });
  
  // Register multiple A record handlers for round-robin DNS
  dnsServer.registerHandler('www.example.com', ['A'], (question) => {
    console.log('First A handler called');
    return {
      name: question.name,
      type: 'A',
      class: 'IN',
      ttl: 300,
      data: '10.0.0.1',
    };
  });
  
  dnsServer.registerHandler('www.example.com', ['A'], (question) => {
    console.log('Second A handler called');
    return {
      name: question.name,
      type: 'A',
      class: 'IN',
      ttl: 300,
      data: '10.0.0.2',
    };
  });
  
  dnsServer.registerHandler('www.example.com', ['A'], (question) => {
    console.log('Third A handler called');
    return {
      name: question.name,
      type: 'A',
      class: 'IN',
      ttl: 300,
      data: '10.0.0.3',
    };
  });
  
  await dnsServer.start();
  
  const client = dgram.createSocket('udp4');
  
  const query = dnsPacket.encode({
    type: 'query',
    id: 2,
    flags: dnsPacket.RECURSION_DESIRED,
    questions: [
      {
        name: 'www.example.com',
        type: 'A',
        class: 'IN',
      },
    ],
  });
  
  const responsePromise = new Promise<dnsPacket.Packet>((resolve, reject) => {
    client.on('message', (msg) => {
      const dnsResponse = dnsPacket.decode(msg);
      resolve(dnsResponse);
      client.close();
    });
    
    client.on('error', (err) => {
      reject(err);
      client.close();
    });
    
    client.send(query, udpPort, 'localhost', (err) => {
      if (err) {
        reject(err);
        client.close();
      }
    });
  });
  
  const dnsResponse = await responsePromise;
  
  console.log('Current behavior - A records returned:', dnsResponse.answers.length);
  console.log('A records:', dnsResponse.answers.map(a => (a as any).data));
  
  // CURRENT BEHAVIOR: Only returns 1 A record, preventing round-robin DNS
  expect(dnsResponse.answers.length).toEqual(1);
  expect((dnsResponse.answers[0] as any).data).toEqual('10.0.0.1');
  
  await stopServer(dnsServer);
  dnsServer = null;
});

tap.test('should demonstrate the limitation with multiple TXT records', async () => {
  const httpsData = await tapNodeTools.createHttpsCert();
  const udpPort = getUniqueUdpPort();
  
  dnsServer = new smartdns.DnsServer({
    httpsKey: httpsData.key,
    httpsCert: httpsData.cert,
    httpsPort: getUniqueHttpsPort(),
    udpPort: udpPort,
    dnssecZone: 'example.com',
  });
  
  // Register multiple TXT record handlers
  dnsServer.registerHandler('example.com', ['TXT'], (question) => {
    console.log('SPF handler called');
    return {
      name: question.name,
      type: 'TXT',
      class: 'IN',
      ttl: 3600,
      data: ['v=spf1 include:_spf.example.com ~all'],
    };
  });
  
  dnsServer.registerHandler('example.com', ['TXT'], (question) => {
    console.log('DKIM handler called');
    return {
      name: question.name,
      type: 'TXT',
      class: 'IN',
      ttl: 3600,
      data: ['v=DKIM1; k=rsa; p=MIGfMA0GCSqGSIb3DQEBAQUAA4GNA...'],
    };
  });
  
  dnsServer.registerHandler('example.com', ['TXT'], (question) => {
    console.log('Domain verification handler called');
    return {
      name: question.name,
      type: 'TXT',
      class: 'IN',
      ttl: 3600,
      data: ['google-site-verification=1234567890abcdef'],
    };
  });
  
  await dnsServer.start();
  
  const client = dgram.createSocket('udp4');
  
  const query = dnsPacket.encode({
    type: 'query',
    id: 3,
    flags: dnsPacket.RECURSION_DESIRED,
    questions: [
      {
        name: 'example.com',
        type: 'TXT',
        class: 'IN',
      },
    ],
  });
  
  const responsePromise = new Promise<dnsPacket.Packet>((resolve, reject) => {
    client.on('message', (msg) => {
      const dnsResponse = dnsPacket.decode(msg);
      resolve(dnsResponse);
      client.close();
    });
    
    client.on('error', (err) => {
      reject(err);
      client.close();
    });
    
    client.send(query, udpPort, 'localhost', (err) => {
      if (err) {
        reject(err);
        client.close();
      }
    });
  });
  
  const dnsResponse = await responsePromise;
  
  console.log('Current behavior - TXT records returned:', dnsResponse.answers.length);
  console.log('TXT records:', dnsResponse.answers.map(a => (a as any).data));
  
  // CURRENT BEHAVIOR: Only returns 1 TXT record instead of all 3
  expect(dnsResponse.answers.length).toEqual(1);
  expect((dnsResponse.answers[0] as any).data[0]).toInclude('spf1');
  
  await stopServer(dnsServer);
  dnsServer = null;
});

tap.test('should show the current workaround pattern', async () => {
  const httpsData = await tapNodeTools.createHttpsCert();
  const udpPort = getUniqueUdpPort();
  
  dnsServer = new smartdns.DnsServer({
    httpsKey: httpsData.key,
    httpsCert: httpsData.cert,
    httpsPort: getUniqueHttpsPort(),
    udpPort: udpPort,
    dnssecZone: 'example.com',
  });
  
  // WORKAROUND: Create an array to store NS records and return them from a single handler
  const nsRecords = ['ns1.example.com', 'ns2.example.com'];
  let nsIndex = 0;
  
  // This workaround still doesn't solve the problem because only one handler executes
  dnsServer.registerHandler('example.com', ['NS'], (question) => {
    const record = nsRecords[nsIndex % nsRecords.length];
    nsIndex++;
    return {
      name: question.name,
      type: 'NS',
      class: 'IN',
      ttl: 3600,
      data: record,
    };
  });
  
  await dnsServer.start();
  
  // Make two queries to show the workaround behavior
  const client1 = dgram.createSocket('udp4');
  const client2 = dgram.createSocket('udp4');
  
  const query = dnsPacket.encode({
    type: 'query',
    id: 4,
    flags: dnsPacket.RECURSION_DESIRED,
    questions: [
      {
        name: 'example.com',
        type: 'NS',
        class: 'IN',
      },
    ],
  });
  
  const responsePromise1 = new Promise<dnsPacket.Packet>((resolve, reject) => {
    client1.on('message', (msg) => {
      const dnsResponse = dnsPacket.decode(msg);
      resolve(dnsResponse);
      client1.close();
    });
    
    client1.send(query, udpPort, 'localhost');
  });
  
  const responsePromise2 = new Promise<dnsPacket.Packet>((resolve, reject) => {
    client2.on('message', (msg) => {
      const dnsResponse = dnsPacket.decode(msg);
      resolve(dnsResponse);
      client2.close();
    });
    
    setTimeout(() => {
      client2.send(query, udpPort, 'localhost');
    }, 100);
  });
  
  const [response1, response2] = await Promise.all([responsePromise1, responsePromise2]);
  
  console.log('First query NS:', (response1.answers[0] as any).data);
  console.log('Second query NS:', (response2.answers[0] as any).data);
  
  // This workaround rotates between records but still only returns one at a time
  expect(response1.answers.length).toEqual(1);
  expect(response2.answers.length).toEqual(1);
  expect((response1.answers[0] as any).data).toEqual('ns1.example.com');
  expect((response2.answers[0] as any).data).toEqual('ns2.example.com');
  
  await stopServer(dnsServer);
  dnsServer = null;
});

export default tap.start();