import * as plugins from './plugins.js';
import { DnsSec } from './classes.dnssec.js';
import * as dnsPacket from 'dns-packet';

interface IDnsServerOptions {
  httpsKey: string;
  httpsCert: string;
  httpsPort: number;
  udpPort: number;
  dnssecZone: string;
}

interface DnsAnswer {
  name: string;
  type: string;
  class: string | number;
  ttl: number;
  data: any;
}

interface IDnsHandler {
  domainPattern: string;
  recordTypes: string[];
  handler: (question: dnsPacket.Question) => DnsAnswer | null;
}

// Define types for DNSSEC records if not provided
interface DNSKEYData {
  flags: number;
  algorithm: number;
  key: Buffer;
}

interface RRSIGData {
  typeCovered: string; // Changed to string to match dns-packet expectations
  algorithm: number;
  labels: number;
  originalTTL: number;
  expiration: number;
  inception: number;
  keyTag: number;
  signerName: string;
  signature: Buffer;
}

export class DnsServer {
  private udpServer: plugins.dgram.Socket;
  private httpsServer: plugins.https.Server;
  private handlers: IDnsHandler[] = [];

  // DNSSEC related properties
  private dnsSec: DnsSec;
  private dnskeyRecord: DNSKEYData;
  private keyTag: number;

  constructor(private options: IDnsServerOptions) {
    // Initialize DNSSEC
    this.dnsSec = new DnsSec({
      zone: options.dnssecZone,
      algorithm: 'ECDSA', // You can change this based on your needs
      keySize: 256,
      days: 365,
    });

    // Generate DNSKEY and DS records
    const { dsRecord, dnskeyRecord } = this.dnsSec.getDsAndKeyPair();

    // Parse DNSKEY record into dns-packet format
    this.dnskeyRecord = this.parseDNSKEYRecord(dnskeyRecord);
    this.keyTag = this.computeKeyTag(this.dnskeyRecord);
  }

  public registerHandler(
    domainPattern: string,
    recordTypes: string[],
    handler: (question: dnsPacket.Question) => DnsAnswer | null
  ): void {
    this.handlers.push({ domainPattern, recordTypes, handler });
  }

  private processDnsRequest(request: dnsPacket.Packet): dnsPacket.Packet {
    const response: dnsPacket.Packet = {
      type: 'response',
      id: request.id,
      flags:
        dnsPacket.AUTHORITATIVE_ANSWER |
        dnsPacket.RECURSION_AVAILABLE |
        (request.flags & dnsPacket.RECURSION_DESIRED ? dnsPacket.RECURSION_DESIRED : 0),
      questions: request.questions,
      answers: [],
      additionals: [],
    };

    const dnssecRequested = this.isDnssecRequested(request);

    for (const question of request.questions) {
      console.log(`Query for ${question.name} of type ${question.type}`);

      let answered = false;

      // Handle DNSKEY queries if DNSSEC is requested
      if (dnssecRequested && question.type === 'DNSKEY' && question.name === this.options.dnssecZone) {
        const dnskeyAnswer: DnsAnswer = {
          name: question.name,
          type: 'DNSKEY',
          class: 'IN',
          ttl: 3600,
          data: this.dnskeyRecord,
        };
        response.answers.push(dnskeyAnswer as plugins.dnsPacket.Answer);

        // Sign the DNSKEY RRset
        const rrsig = this.generateRRSIG('DNSKEY', [dnskeyAnswer], question.name);
        response.answers.push(rrsig as plugins.dnsPacket.Answer);

        answered = true;
        continue;
      }

      for (const handlerEntry of this.handlers) {
        if (
          plugins.minimatch.minimatch(question.name, handlerEntry.domainPattern) &&
          handlerEntry.recordTypes.includes(question.type)
        ) {
          const answer = handlerEntry.handler(question);
          if (answer) {
            // Ensure the answer has ttl and class
            const dnsAnswer: DnsAnswer = {
              ...answer,
              ttl: answer.ttl || 300,
              class: answer.class || 'IN',
            };
            response.answers.push(dnsAnswer as plugins.dnsPacket.Answer);

            if (dnssecRequested) {
              // Sign the answer RRset
              const rrsig = this.generateRRSIG(question.type, [dnsAnswer], question.name);
              response.answers.push(rrsig as plugins.dnsPacket.Answer);
            }

            answered = true;
            break;
          }
        }
      }

      if (!answered) {
        console.log(`No handler found for ${question.name} of type ${question.type}`);
        response.flags |= dnsPacket.AUTHORITATIVE_ANSWER;
        const soaAnswer: DnsAnswer = {
          name: question.name,
          type: 'SOA',
          class: 'IN',
          ttl: 3600,
          data: {
            mname: `ns1.${this.options.dnssecZone}`,
            rname: `hostmaster.${this.options.dnssecZone}`,
            serial: Math.floor(Date.now() / 1000),
            refresh: 3600,
            retry: 600,
            expire: 604800,
            minimum: 86400,
          },
        };
        response.answers.push(soaAnswer as plugins.dnsPacket.Answer);
      }
    }

    return response;
  }

  private isDnssecRequested(request: dnsPacket.Packet): boolean {
    if (!request.additionals) return false;
    for (const additional of request.additionals) {
      if (additional.type === 'OPT' && typeof additional.flags === 'number') {
        // The DO bit is the 15th bit (0x8000)
        if (additional.flags & 0x8000) {
          return true;
        }
      }
    }
    return false;
  }

  private generateRRSIG(
    type: string,
    rrset: DnsAnswer[],
    name: string
  ): DnsAnswer {
    // Prepare RRSIG data
    const algorithm = this.dnsSec.getAlgorithmNumber();
    const keyTag = this.keyTag;
    const signerName = this.options.dnssecZone.endsWith('.') ? this.options.dnssecZone : `${this.options.dnssecZone}.`;
    const inception = Math.floor(Date.now() / 1000) - 3600; // 1 hour ago
    const expiration = inception + 86400; // Valid for 1 day
    const ttl = rrset[0].ttl || 300;

    // Serialize the RRset in canonical form
    const rrsetBuffer = this.serializeRRset(rrset);

    // Sign the RRset
    const signature = this.dnsSec.signData(rrsetBuffer);

    // Construct the RRSIG record
    const rrsig: DnsAnswer = {
      name,
      type: 'RRSIG',
      class: 'IN',
      ttl,
      data: {
        typeCovered: type, // Changed to type string
        algorithm,
        labels: name.split('.').length - 1,
        originalTTL: ttl,
        expiration,
        inception,
        keyTag,
        signerName,
        signature: signature,
      },
    };

    return rrsig;
  }

  private serializeRRset(rrset: DnsAnswer[]): Buffer {
    // Implement canonical DNS RRset serialization as per RFC 4034 Section 6
    const buffers: Buffer[] = [];
    for (const rr of rrset) {
      if (rr.type === 'OPT') {
        continue; // Skip OPT records
      }

      const name = rr.name.endsWith('.') ? rr.name : rr.name + '.';
      const nameBuffer = this.nameToBuffer(name.toLowerCase());

      const typeValue = this.qtypeToNumber(rr.type);
      const typeBuffer = Buffer.alloc(2);
      typeBuffer.writeUInt16BE(typeValue, 0);

      const classValue = this.classToNumber(rr.class);
      const classBuffer = Buffer.alloc(2);
      classBuffer.writeUInt16BE(classValue, 0);

      const ttlValue = rr.ttl || 300;
      const ttlBuffer = Buffer.alloc(4);
      ttlBuffer.writeUInt32BE(ttlValue, 0);

      // Serialize the data based on the record type
      const dataBuffer = this.serializeRData(rr.type, rr.data);

      const rdLengthBuffer = Buffer.alloc(2);
      rdLengthBuffer.writeUInt16BE(dataBuffer.length, 0);

      buffers.push(Buffer.concat([nameBuffer, typeBuffer, classBuffer, ttlBuffer, rdLengthBuffer, dataBuffer]));
    }
    return Buffer.concat(buffers);
  }

  private serializeRData(type: string, data: any): Buffer {
    // Implement serialization for each record type you support
    switch (type) {
      case 'A':
        return Buffer.from(data.split('.').map((octet: string) => parseInt(octet, 10)));
      case 'AAAA':
        // Handle IPv6 addresses
        return Buffer.from(data.split(':').flatMap((segment: string) => {
          const num = parseInt(segment, 16);
          return [num >> 8, num & 0xff];
        }));
      case 'DNSKEY':
        const dnskeyData: DNSKEYData = data;
        return Buffer.concat([
          Buffer.from([dnskeyData.flags >> 8, dnskeyData.flags & 0xff]),
          Buffer.from([3]), // Protocol field, always 3
          Buffer.from([dnskeyData.algorithm]),
          dnskeyData.key,
        ]);
      case 'SOA':
        // Implement SOA record serialization if needed
        // For now, return an empty buffer or handle as needed
        return Buffer.alloc(0);
      // Add cases for other record types as needed
      default:
        throw new Error(`Serialization for record type ${type} is not implemented.`);
    }
  }

  private parseDNSKEYRecord(dnskeyRecord: string): DNSKEYData {
    // Parse the DNSKEY record string into dns-packet format
    const parts = dnskeyRecord.trim().split(/\s+/);
    const flags = parseInt(parts[3], 10);
    const algorithm = parseInt(parts[5], 10);
    const publicKeyBase64 = parts.slice(6).join('');
    const key = Buffer.from(publicKeyBase64, 'base64');

    return {
      flags,
      algorithm,
      key,
    };
  }

  private computeKeyTag(dnskeyRecord: DNSKEYData): number {
    // Compute key tag as per RFC 4034 Appendix B
    const flags = dnskeyRecord.flags;
    const algorithm = dnskeyRecord.algorithm;
    const key = dnskeyRecord.key;

    const dnskeyRdata = Buffer.concat([
      Buffer.from([flags >> 8, flags & 0xff]),
      Buffer.from([3]), // Protocol field, always 3
      Buffer.from([algorithm]),
      key,
    ]);

    let acc = 0;
    for (let i = 0; i < dnskeyRdata.length; i++) {
      acc += (i & 1) ? dnskeyRdata[i] : dnskeyRdata[i] << 8;
    }
    acc += (acc >> 16) & 0xffff;
    return acc & 0xffff;
  }

  private handleHttpsRequest(req: plugins.http.IncomingMessage, res: plugins.http.ServerResponse): void {
    if (req.method === 'POST' && req.url === '/dns-query') {
      let body: Buffer[] = [];

      req.on('data', (chunk) => {
        body.push(chunk);
      }).on('end', () => {
        const msg = Buffer.concat(body);
        const request = dnsPacket.decode(msg);
        const response = this.processDnsRequest(request);
        const responseData = dnsPacket.encode(response);
        res.writeHead(200, { 'Content-Type': 'application/dns-message' });
        res.end(responseData);
      });
    } else {
      res.writeHead(404);
      res.end();
    }
  }

  public async start(): Promise<void> {
    this.httpsServer = plugins.https.createServer(
      {
        key: this.options.httpsKey,
        cert: this.options.httpsCert,
      },
      this.handleHttpsRequest.bind(this)
    );

    this.udpServer = plugins.dgram.createSocket('udp4');
    this.udpServer.on('message', (msg, rinfo) => {
      const request = dnsPacket.decode(msg);
      const response = this.processDnsRequest(request);
      const responseData = dnsPacket.encode(response);
      this.udpServer.send(responseData, rinfo.port, rinfo.address);
    });

    this.udpServer.on('error', (err) => {
      console.error(`UDP Server error:\n${err.stack}`);
      this.udpServer.close();
    });

    const udpListeningDeferred = plugins.smartpromise.defer<void>();
    const httpsListeningDeferred = plugins.smartpromise.defer<void>();

    try {
      this.udpServer.bind(this.options.udpPort, '0.0.0.0', () => {
        console.log(`UDP DNS server running on port ${this.options.udpPort}`);
        udpListeningDeferred.resolve();
      });

      this.httpsServer.listen(this.options.httpsPort, () => {
        console.log(`HTTPS DNS server running on port ${this.options.httpsPort}`);
        httpsListeningDeferred.resolve();
      });
    } catch (err) {
      console.error('Error starting DNS server:', err);
      process.exit(1);
    }
    await Promise.all([udpListeningDeferred.promise, httpsListeningDeferred.promise]);
  }

  public async stop(): Promise<void> {
    const doneUdp = plugins.smartpromise.defer<void>();
    const doneHttps = plugins.smartpromise.defer<void>();
    this.udpServer.close(() => {
      console.log('UDP DNS server stopped');
      this.udpServer.unref();
      this.udpServer = null;
      doneUdp.resolve();
    });

    this.httpsServer.close(() => {
      console.log('HTTPS DNS server stopped');
      this.httpsServer.unref();
      this.httpsServer = null;
      doneHttps.resolve();
    });

    await Promise.all([doneUdp.promise, doneHttps.promise]);
  }

  // Helper methods

  private qtypeToNumber(type: string): number {
    const QTYPE_NUMBERS: { [key: string]: number } = {
      'A': 1,
      'NS': 2,
      'CNAME': 5,
      'SOA': 6,
      'PTR': 12,
      'MX': 15,
      'TXT': 16,
      'AAAA': 28,
      'SRV': 33,
      'DNSKEY': 48,
      'RRSIG': 46,
      // Add more as needed
    };
    return QTYPE_NUMBERS[type.toUpperCase()] || 0;
  }

  private classToNumber(cls: string | number): number {
    const CLASS_NUMBERS: { [key: string]: number } = {
      'IN': 1,
      'CH': 3,
      'HS': 4,
      // Add more as needed
    };
    if (typeof cls === 'number') {
      return cls;
    }
    return CLASS_NUMBERS[cls.toUpperCase()] || 1;
  }

  private nameToBuffer(name: string): Buffer {
    const labels = name.split('.');
    const buffers = labels.map(label => {
      const len = Buffer.byteLength(label, 'utf8');
      const buf = Buffer.alloc(1 + len);
      buf.writeUInt8(len, 0);
      buf.write(label, 1);
      return buf;
    });
    return Buffer.concat([...buffers, Buffer.from([0])]); // Add root label
  }
}