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

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

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

export 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;
}

// Let's Encrypt related interfaces
interface LetsEncryptOptions {
  email?: string;
  staging?: boolean;
  certDir?: string;
}

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 });
  }

  // Unregister a specific handler
  public unregisterHandler(domainPattern: string, recordTypes: string[]): boolean {
    const initialLength = this.handlers.length;
    this.handlers = this.handlers.filter(handler => 
      !(handler.domainPattern === domainPattern && 
        recordTypes.every(type => handler.recordTypes.includes(type)))
    );
    return this.handlers.length < initialLength;
  }

  /**
   * Retrieve SSL certificate for specified domains using Let's Encrypt
   * @param domainNames Array of domain names to include in the certificate
   * @param options Configuration options for Let's Encrypt
   * @returns Object containing certificate, private key, and success status
   */
  public async retrieveSslCertificate(
    domainNames: string[],
    options: LetsEncryptOptions = {}
  ): Promise<{ cert: string; key: string; success: boolean }> {
    // Default options
    const opts = {
      email: options.email || 'admin@example.com',
      staging: options.staging !== undefined ? options.staging : false,
      certDir: options.certDir || './certs'
    };

    // Create certificate directory if it doesn't exist
    if (!plugins.fs.existsSync(opts.certDir)) {
      plugins.fs.mkdirSync(opts.certDir, { recursive: true });
    }

    // Filter domains this server is authoritative for
    const authorizedDomains = this.filterAuthorizedDomains(domainNames);
    
    if (authorizedDomains.length === 0) {
      console.error('None of the provided domains are authorized for this DNS server');
      return { cert: '', key: '', success: false };
    }
    
    console.log(`Retrieving SSL certificate for domains: ${authorizedDomains.join(', ')}`);
    
    try {
      // Allow for override in tests
      // @ts-ignore - acmeClientOverride is added for testing purposes
      const acmeClient = this.acmeClientOverride || await import('acme-client');
      
      // Generate or load account key
      const accountKeyPath = plugins.path.join(opts.certDir, 'account.key');
      let accountKey: Buffer;
      
      if (plugins.fs.existsSync(accountKeyPath)) {
        accountKey = plugins.fs.readFileSync(accountKeyPath);
      } else {
        // Generate new account key
        const { privateKey } = plugins.crypto.generateKeyPairSync('rsa', {
          modulusLength: 2048,
          publicKeyEncoding: { type: 'spki', format: 'pem' },
          privateKeyEncoding: { type: 'pkcs8', format: 'pem' }
        });
        
        accountKey = Buffer.from(privateKey);
        plugins.fs.writeFileSync(accountKeyPath, accountKey);
      }
      
      // Initialize ACME client
      const client = new acmeClient.Client({
        directoryUrl: opts.staging
          ? acmeClient.directory.letsencrypt.staging
          : acmeClient.directory.letsencrypt.production,
        accountKey: accountKey
      });
      
      // Create or update account
      await client.createAccount({
        termsOfServiceAgreed: true,
        contact: [`mailto:${opts.email}`]
      });
      
      // Create order for certificate
      const order = await client.createOrder({
        identifiers: authorizedDomains.map(domain => ({
          type: 'dns',
          value: domain
        }))
      });
      
      // Get authorizations
      const authorizations = await client.getAuthorizations(order);
      
      // Track handlers to clean up later
      const challengeHandlers: { domain: string; pattern: string }[] = [];
      
      // Process each authorization
      for (const auth of authorizations) {
        const domain = auth.identifier.value;
        
        // Get DNS challenge
        const challenge = auth.challenges.find(c => c.type === 'dns-01');
        if (!challenge) {
          throw new Error(`No DNS-01 challenge found for ${domain}`);
        }
        
        // Get key authorization and DNS record value
        const keyAuthorization = await client.getChallengeKeyAuthorization(challenge);
        const recordValue = this.getDnsRecordValueForChallenge(keyAuthorization);
        
        // Create challenge domain (where TXT record should be placed)
        const challengeDomain = `_acme-challenge.${domain}`;
        
        console.log(`Setting up TXT record for ${challengeDomain}: ${recordValue}`);
        
        // Register handler for the TXT record
        this.registerHandler(
          challengeDomain,
          ['TXT'],
          (question: dnsPacket.Question): DnsAnswer | null => {
            if (question.name === challengeDomain && question.type === 'TXT') {
              return {
                name: question.name,
                type: 'TXT',
                class: 'IN',
                ttl: 300,
                data: [recordValue]
              };
            }
            return null;
          }
        );
        
        // Track the handler for cleanup
        challengeHandlers.push({ domain, pattern: challengeDomain });
        
        // Wait briefly for DNS propagation
        await new Promise(resolve => setTimeout(resolve, 2000));
        
        // Complete the challenge
        await client.completeChallenge(challenge);
        
        // Wait for verification
        await client.waitForValidStatus(challenge);
        console.log(`Challenge for ${domain} validated successfully!`);
      }
      
      // Generate certificate key
      const domainKeyPath = plugins.path.join(opts.certDir, `${authorizedDomains[0]}.key`);
      const { privateKey } = plugins.crypto.generateKeyPairSync('rsa', {
        modulusLength: 2048,
        publicKeyEncoding: { type: 'spki', format: 'pem' },
        privateKeyEncoding: { type: 'pkcs8', format: 'pem' }
      });
      
      plugins.fs.writeFileSync(domainKeyPath, privateKey);
      
      // Create CSR
      // Define an interface for the expected CSR result structure
      interface CSRResult {
        csr: Buffer;
      }
      
      // Use the forge.createCsr method and handle typing with a more direct approach
      const csrResult = await acmeClient.forge.createCsr({
        commonName: authorizedDomains[0],
        altNames: authorizedDomains
      }) as unknown as CSRResult;
      
      // Finalize the order with the CSR
      await client.finalizeOrder(order, csrResult.csr);
      
      // Get certificate
      const certificate = await client.getCertificate(order);
      
      // Save certificate
      const certPath = plugins.path.join(opts.certDir, `${authorizedDomains[0]}.cert`);
      plugins.fs.writeFileSync(certPath, certificate);
      
      // Update HTTPS server with new certificate
      this.options.httpsCert = certificate;
      this.options.httpsKey = privateKey;
      
      // Restart HTTPS server with new certificate
      await this.restartHttpsServer();
      
      // Clean up challenge handlers
      for (const handler of challengeHandlers) {
        this.unregisterHandler(handler.pattern, ['TXT']);
        console.log(`Cleaned up challenge handler for ${handler.domain}`);
      }
      
      return {
        cert: certificate,
        key: privateKey,
        success: true
      };
    } catch (error) {
      console.error('Error retrieving SSL certificate:', error);
      return { cert: '', key: '', success: false };
    }
  }

  /**
   * Create DNS record value for the ACME challenge
   */
  private getDnsRecordValueForChallenge(keyAuthorization: string): string {
    // Create SHA-256 digest of the key authorization
    const digest = plugins.crypto
      .createHash('sha256')
      .update(keyAuthorization)
      .digest('base64')
      .replace(/\+/g, '-')
      .replace(/\//g, '_')
      .replace(/=/g, '');
    
    return digest;
  }

  /**
   * Restart the HTTPS server with the new certificate
   */
  private async restartHttpsServer(): Promise<void> {
    return new Promise<void>((resolve, reject) => {
      // First check if the server exists
      if (!this.httpsServer) {
        console.log('No HTTPS server to restart');
        resolve();
        return;
      }

      this.httpsServer.close(() => {
        try {
          // Validate certificate and key before trying to create the server
          if (!this.options.httpsCert || !this.options.httpsKey) {
            throw new Error('Missing certificate or key for HTTPS server');
          }

          // For testing, check if we have a mock certificate
          if (this.options.httpsCert.includes('MOCK_CERTIFICATE')) {
            console.log('Using mock certificate in test mode');
            // In test mode with mock cert, we can use the original cert
            // @ts-ignore - accessing acmeClientOverride for testing
            if (this.acmeClientOverride) {
              this.httpsServer = plugins.https.createServer(
                {
                  key: this.options.httpsKey,
                  cert: this.options.httpsCert,
                },
                this.handleHttpsRequest.bind(this)
              );
              
              this.httpsServer.listen(this.options.httpsPort, () => {
                console.log(`HTTPS DNS server restarted on port ${this.options.httpsPort} with test certificate`);
                resolve();
              });
              return;
            }
          }

          // Create the new server with the updated certificate
          this.httpsServer = plugins.https.createServer(
            {
              key: this.options.httpsKey,
              cert: this.options.httpsCert,
            },
            this.handleHttpsRequest.bind(this)
          );
          
          this.httpsServer.listen(this.options.httpsPort, () => {
            console.log(`HTTPS DNS server restarted on port ${this.options.httpsPort} with new certificate`);
            resolve();
          });
        } catch (err) {
          console.error('Error creating HTTPS server with new certificate:', err);
          reject(err);
        }
      });
    });
  }

  /**
   * Filter domains to include only those the server is authoritative for
   */
  public filterAuthorizedDomains(domainNames: string[]): string[] {
    const authorizedDomains: string[] = [];
    
    for (const domain of domainNames) {
      // Handle wildcards (*.example.com)
      if (domain.startsWith('*.')) {
        const baseDomain = domain.substring(2);
        if (this.isAuthorizedForDomain(baseDomain)) {
          authorizedDomains.push(domain);
        }
      } 
      // Regular domains
      else if (this.isAuthorizedForDomain(domain)) {
        authorizedDomains.push(domain);
      }
    }
    
    return authorizedDomains;
  }

  /**
   * Check if the server is authoritative for a domain
   */
  private isAuthorizedForDomain(domain: string): boolean {
    // Check if any handler matches this domain
    for (const handler of this.handlers) {
      if (plugins.minimatch.minimatch(domain, handler.domainPattern)) {
        return true;
      }
    }
    
    // Also check if the domain is the DNSSEC zone itself
    if (domain === this.options.dnssecZone || domain.endsWith(`.${this.options.dnssecZone}`)) {
      return true;
    }
    
    return false;
  }

  public 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 'TXT':
        // Handle TXT records for ACME challenges
        if (Array.isArray(data)) {
          // Combine all strings and encode as lengths and values
          const buffers = data.map(str => {
            const strBuf = Buffer.from(str);
            const lenBuf = Buffer.alloc(1);
            lenBuf.writeUInt8(strBuf.length, 0);
            return Buffer.concat([lenBuf, strBuf]);
          });
          return Buffer.concat(buffers);
        }
        return Buffer.alloc(0);
      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>();
    
    if (this.udpServer) {
      this.udpServer.close(() => {
        console.log('UDP DNS server stopped');
        if (this.udpServer) {
          this.udpServer.unref();
          this.udpServer = null;
        }
        doneUdp.resolve();
      });
    } else {
      doneUdp.resolve();
    }

    if (this.httpsServer) {
      this.httpsServer.close(() => {
        console.log('HTTPS DNS server stopped');
        if (this.httpsServer) {
          this.httpsServer.unref();
          this.httpsServer = null;
        }
        doneHttps.resolve();
      });
    } else {
      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
  }
}