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 = 8500;
let nextUdpPort = 8501;

function getUniqueHttpsPort() {
  return nextHttpsPort++;
}

function getUniqueUdpPort() {
  return nextUdpPort++;
}

// Cleanup function for servers
async function stopServer(server: smartdns.DnsServer | null | undefined) {
  if (!server) {
    return;
  }
  
  try {
    await server.stop();
  } catch (e) {
    console.log('Handled error when stopping server:', e.message || e);
  }
}

tap.test('DNSSEC should sign entire RRset together, not individual 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
  dnsServer.registerHandler('example.com', ['NS'], (question) => {
    return {
      name: question.name,
      type: 'NS',
      class: 'IN',
      ttl: 3600,
      data: 'ns1.example.com',
    };
  });
  
  dnsServer.registerHandler('example.com', ['NS'], (question) => {
    return {
      name: question.name,
      type: 'NS',
      class: 'IN',
      ttl: 3600,
      data: 'ns2.example.com',
    };
  });
  
  dnsServer.registerHandler('example.com', ['NS'], (question) => {
    return {
      name: question.name,
      type: 'NS',
      class: 'IN',
      ttl: 3600,
      data: 'ns3.example.com',
    };
  });
  
  await dnsServer.start();
  
  const client = dgram.createSocket('udp4');
  
  // Create query with DNSSEC requested
  const query = dnsPacket.encode({
    type: 'query',
    id: 1,
    flags: dnsPacket.RECURSION_DESIRED,
    questions: [
      {
        name: 'example.com',
        type: 'NS',
        class: 'IN',
      },
    ],
    additionals: [
      {
        name: '.',
        type: 'OPT',
        ttl: 0,
        flags: 0x8000, // DO bit set for DNSSEC
        data: Buffer.alloc(0),
      } as any,
    ],
  });
  
  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;
  
  // Count NS and RRSIG records
  const nsAnswers = dnsResponse.answers.filter(a => a.type === 'NS');
  const rrsigAnswers = dnsResponse.answers.filter(a => a.type === 'RRSIG');
  
  console.log('NS records returned:', nsAnswers.length);
  console.log('RRSIG records returned:', rrsigAnswers.length);
  
  // Should have 3 NS records and only 1 RRSIG for the entire RRset
  expect(nsAnswers.length).toEqual(3);
  expect(rrsigAnswers.length).toEqual(1);
  
  // Verify RRSIG covers NS type
  const rrsigData = (rrsigAnswers[0] as any).data;
  expect(rrsigData.typeCovered).toEqual('NS');
  
  await stopServer(dnsServer);
  dnsServer = null;
});

tap.test('SOA records should be properly serialized and returned', 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',
  });
  
  await dnsServer.start();
  
  const client = dgram.createSocket('udp4');
  
  // Query for a non-existent subdomain to trigger SOA response
  const query = dnsPacket.encode({
    type: 'query',
    id: 2,
    flags: dnsPacket.RECURSION_DESIRED,
    questions: [
      {
        name: 'nonexistent.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;
  
  // Should have SOA record in response
  const soaAnswers = dnsResponse.answers.filter(a => a.type === 'SOA');
  expect(soaAnswers.length).toEqual(1);
  
  const soaData = (soaAnswers[0] as any).data;
  console.log('SOA record:', soaData);
  
  expect(soaData.mname).toEqual('ns1.example.com');
  expect(soaData.rname).toEqual('hostmaster.example.com');
  expect(typeof soaData.serial).toEqual('number');
  expect(soaData.refresh).toEqual(3600);
  expect(soaData.retry).toEqual(600);
  expect(soaData.expire).toEqual(604800);
  expect(soaData.minimum).toEqual(86400);
  
  await stopServer(dnsServer);
  dnsServer = null;
});

tap.test('Primary nameserver should be configurable', 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',
    primaryNameserver: 'custom-ns.example.com',
  });
  
  await dnsServer.start();
  
  const client = dgram.createSocket('udp4');
  
  // Query for SOA record
  const query = dnsPacket.encode({
    type: 'query',
    id: 3,
    flags: dnsPacket.RECURSION_DESIRED,
    questions: [
      {
        name: 'example.com',
        type: 'SOA',
        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;
  
  // Should have SOA record with custom nameserver
  const soaAnswers = dnsResponse.answers.filter(a => a.type === 'SOA');
  expect(soaAnswers.length).toEqual(1);
  
  const soaData = (soaAnswers[0] as any).data;
  console.log('SOA mname:', soaData.mname);
  
  // Should use the custom primary nameserver
  expect(soaData.mname).toEqual('custom-ns.example.com');
  
  await stopServer(dnsServer);
  dnsServer = null;
});

tap.test('Multiple A records should have single RRSIG when DNSSEC is enabled', 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 records for round-robin
  const ips = ['10.0.0.1', '10.0.0.2', '10.0.0.3'];
  for (const ip of ips) {
    dnsServer.registerHandler('www.example.com', ['A'], (question) => {
      return {
        name: question.name,
        type: 'A',
        class: 'IN',
        ttl: 300,
        data: ip,
      };
    });
  }
  
  await dnsServer.start();
  
  const client = dgram.createSocket('udp4');
  
  const query = dnsPacket.encode({
    type: 'query',
    id: 4,
    flags: dnsPacket.RECURSION_DESIRED,
    questions: [
      {
        name: 'www.example.com',
        type: 'A',
        class: 'IN',
      },
    ],
    additionals: [
      {
        name: '.',
        type: 'OPT',
        ttl: 0,
        flags: 0x8000, // DO bit set for DNSSEC
        data: Buffer.alloc(0),
      } as any,
    ],
  });
  
  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;
  
  const aAnswers = dnsResponse.answers.filter(a => a.type === 'A');
  const rrsigAnswers = dnsResponse.answers.filter(a => a.type === 'RRSIG');
  
  console.log('A records:', aAnswers.length);
  console.log('RRSIG records:', rrsigAnswers.length);
  
  // Should have 3 A records and only 1 RRSIG
  expect(aAnswers.length).toEqual(3);
  expect(rrsigAnswers.length).toEqual(1);
  
  await stopServer(dnsServer);
  dnsServer = null;
});

export default tap.start();