const OPCODE_MASK = 0x7800;
const RCODE_MASK = 0x000f;

const FLAG_AA = 0x0400;
const FLAG_TC = 0x0200;
const FLAG_RD = 0x0100;
const FLAG_RA = 0x0080;
const FLAG_AD = 0x0020;

const RRTYPE_A = 1;
const RRTYPE_NS = 2;
const RRTYPE_CNAME = 5;
const RRTYPE_SOA = 6;
const RRTYPE_PTR = 12;
const RRTYPE_MX = 15;
const RRTYPE_TXT = 16;
const RRTYPE_AAAA = 28;
const RRTYPE_SRV = 33;
const RRTYPE_CAA = 257;

const RCODE_NAMES: Record<number, string> = {
  0: "NOERROR",
  1: "FORMERR",
  2: "SERVFAIL",
  3: "NXDOMAIN",
  4: "NOTIMP",
  5: "REFUSED",
  6: "YXDOMAIN",
  7: "YXRRSET",
  8: "NXRRSET",
  9: "NOTAUTH",
  10: "NOTZONE",
};

const RRTYPE_NAMES: Record<number, string> = {
  [RRTYPE_A]: "A",
  [RRTYPE_AAAA]: "AAAA",
  [RRTYPE_CAA]: "CAA",
  [RRTYPE_CNAME]: "CNAME",
  [RRTYPE_MX]: "MX",
  [RRTYPE_NS]: "NS",
  [RRTYPE_PTR]: "PTR",
  [RRTYPE_SOA]: "SOA",
  [RRTYPE_SRV]: "SRV",
  [RRTYPE_TXT]: "TXT",
};

const RRTYPE_BY_NAME: Record<string, number> = {
  A: RRTYPE_A,
  AAAA: RRTYPE_AAAA,
  CAA: RRTYPE_CAA,
  CNAME: RRTYPE_CNAME,
  MX: RRTYPE_MX,
  NS: RRTYPE_NS,
  PTR: RRTYPE_PTR,
  SOA: RRTYPE_SOA,
  SRV: RRTYPE_SRV,
  TXT: RRTYPE_TXT,
};

const CLASS_IN = 1;

export interface DnsAnswer {
  class: number;
  data: Record<string, unknown>;
  name: string;
  ttl: number;
  type: number;
  value: string;
}

export interface DnsFlags {
  authenticatedData: boolean;
  authoritative: boolean;
  recursionAvailable: boolean;
  recursionDesired: boolean;
  truncated: boolean;
}

export interface DnsHeader {
  additionalCount: number;
  answerCount: number;
  authorityCount: number;
  flags: DnsFlags;
  id: number;
  opcode: number;
  questionCount: number;
  rcode: number;
}

export interface DnsQuestion {
  name: string;
  qclass: number;
  qtype: number;
}

export interface DnsResponse {
  additional: DnsAnswer[];
  answers: DnsAnswer[];
  authorities: DnsAnswer[];
  header: DnsHeader;
  questions: DnsQuestion[];
}

interface ParseContext {
  data: Uint8Array;
  offset: number;
  view: DataView;
}

interface RdataResult extends Record<string, unknown> {
  value: string;
}

export function buildQuery(name: string, recordType: number, recursionDesired: boolean): Uint8Array {
  const queryName = name.endsWith(".") ? name.slice(0, -1) : name;
  const labels = queryName === "" ? [] : queryName.split(".");
  const encodedLabels: Uint8Array[] = [];
  let nameSize = 1;

  for (const label of labels) {
    if (label === "") {
      throw new Error(`无效的 DNS 名称: ${name}`);
    }
    if (!isAscii(label)) {
      throw new Error(`DNS 名称必须使用 ASCII/Punycode: ${name}`);
    }

    const encoded = new TextEncoder().encode(label);
    if (encoded.length > 63) {
      throw new Error(`DNS 标签超过 63 字节: ${label}`);
    }
    nameSize += 1 + encoded.length;
    encodedLabels.push(encoded);
  }

  if (nameSize > 255) {
    throw new Error(`DNS 名称超过 255 字节: ${name}`);
  }

  const questionSize = nameSize + 4;

  const buf = new Uint8Array(12 + questionSize);
  const view = new DataView(buf.buffer);
  const id = (Math.random() * 65536) | 0;
  view.setUint16(0, id);

  let flags = 0;
  if (recursionDesired) flags |= FLAG_RD;
  view.setUint16(2, flags);
  view.setUint16(4, 1);
  view.setUint16(6, 0);
  view.setUint16(8, 0);
  view.setUint16(10, 0);

  let offset = 12;
  for (const encoded of encodedLabels) {
    buf[offset] = encoded.length;
    offset++;
    buf.set(encoded, offset);
    offset += encoded.length;
  }
  buf[offset] = 0;
  offset++;

  view.setUint16(offset, recordType);
  offset += 2;
  view.setUint16(offset, CLASS_IN);

  return buf;
}

export function parseResponse(data: Uint8Array): DnsResponse {
  const view = new DataView(data.buffer, data.byteOffset, data.byteLength);

  if (data.byteLength < 12) {
    throw new Error(`DNS 响应过短: ${data.byteLength} 字节`);
  }

  const id = view.getUint16(0);
  const flagsWord = view.getUint16(2);
  const questionCount = view.getUint16(4);
  const answerCount = view.getUint16(6);
  const authorityCount = view.getUint16(8);
  const additionalCount = view.getUint16(10);

  const opcode = (flagsWord & OPCODE_MASK) >> 11;
  const rcode = flagsWord & RCODE_MASK;

  const flags: DnsFlags = {
    authenticatedData: (flagsWord & FLAG_AD) !== 0,
    authoritative: (flagsWord & FLAG_AA) !== 0,
    recursionAvailable: (flagsWord & FLAG_RA) !== 0,
    recursionDesired: (flagsWord & FLAG_RD) !== 0,
    truncated: (flagsWord & FLAG_TC) !== 0,
  };

  const header: DnsHeader = {
    additionalCount,
    answerCount,
    authorityCount,
    flags,
    id,
    opcode,
    questionCount,
    rcode,
  };

  const ctx: ParseContext = { data, offset: 12, view };
  const questions: DnsQuestion[] = [];
  for (let i = 0; i < questionCount; i++) {
    const name = readName(ctx);
    const qtype = readUint16(ctx);
    const qclass = readUint16(ctx);
    questions.push({ name, qclass, qtype });
  }

  const answers = readAnswers(ctx, answerCount);
  const authorities = readAnswers(ctx, authorityCount);
  const additional = readAnswers(ctx, additionalCount);

  return { additional, answers, authorities, header, questions };
}

export function rcodeName(code: number): string {
  return RCODE_NAMES[code] ?? `UNKNOWN(${code})`;
}

export function rrtTypeByName(name: string): number {
  const code = RRTYPE_BY_NAME[name];
  if (code === undefined) throw new Error(`不支持的记录类型: ${name}`);
  return code;
}

export function rrtTypeName(code: number): string {
  return RRTYPE_NAMES[code] ?? `TYPE${code}`;
}

function isAscii(value: string): boolean {
  for (let i = 0; i < value.length; i++) {
    if (value.charCodeAt(i) > 0x7f) return false;
  }
  return true;
}

function normalizeIPv6(groups: string[]): string {
  let maxZeroStart = -1;
  let maxZeroLen = 0;
  let currentStart = -1;
  let currentLen = 0;

  for (let i = 0; i < 8; i++) {
    if (groups[i] === "0") {
      if (currentStart === -1) currentStart = i;
      currentLen++;
      if (currentLen > maxZeroLen) {
        maxZeroLen = currentLen;
        maxZeroStart = currentStart;
      }
    } else {
      currentStart = -1;
      currentLen = 0;
    }
  }

  if (maxZeroLen >= 2) {
    const before = groups.slice(0, maxZeroStart).join(":");
    const after = groups.slice(maxZeroStart + maxZeroLen).join(":");
    if (before === "" && after === "") return "::";
    if (before === "") return `::${after}`;
    if (after === "") return `${before}::`;
    return `${before}::${after}`;
  }

  return groups.join(":");
}

function parseRdata(ctx: ParseContext, type: number, rdLength: number): RdataResult {
  const rdStart = ctx.offset;
  const rdEnd = rdStart + rdLength;

  switch (type) {
    case RRTYPE_A: {
      if (rdLength < 4) throw new Error("A 记录 rdata 过短");
      const a = ctx.data[rdStart]!;
      const b = ctx.data[rdStart + 1]!;
      const c = ctx.data[rdStart + 2]!;
      const d = ctx.data[rdStart + 3]!;
      const address = `${a}.${b}.${c}.${d}`;
      return { address, value: address };
    }
    case RRTYPE_AAAA: {
      if (rdLength < 16) throw new Error("AAAA 记录 rdata 过短");
      const groups: string[] = [];
      for (let i = 0; i < 16; i += 2) {
        groups.push(((ctx.data[rdStart + i]! << 8) | ctx.data[rdStart + i + 1]!).toString(16));
      }
      const address = normalizeIPv6(groups);
      return { address, value: address };
    }
    case RRTYPE_CAA: {
      const flags = ctx.data[rdStart]!;
      const tagLen = ctx.data[rdStart + 1]!;
      const tag = new TextDecoder().decode(ctx.data.subarray(rdStart + 2, rdStart + 2 + tagLen));
      const valueBytes = ctx.data.subarray(rdStart + 2 + tagLen, rdEnd);
      const caaValue = new TextDecoder().decode(valueBytes);
      const value = `${flags} ${tag} ${caaValue}`;
      return { flags, tag, value: value, valueStr: caaValue };
    }
    case RRTYPE_CNAME: {
      const target = readName(ctx);
      return { target, value: target };
    }
    case RRTYPE_MX: {
      const preference = readUint16(ctx);
      const exchange = readName(ctx);
      const value = `${preference} ${exchange}`;
      return { exchange, preference, value };
    }
    case RRTYPE_NS: {
      const nsdname = readName(ctx);
      return { nsdname, value: nsdname };
    }
    case RRTYPE_PTR: {
      const ptrdname = readName(ctx);
      return { ptrdname, value: ptrdname };
    }
    case RRTYPE_SOA: {
      const mname = readName(ctx);
      const rname = readName(ctx);
      const serial = readUint32(ctx);
      const refresh = readUint32(ctx);
      const retry = readUint32(ctx);
      const expire = readUint32(ctx);
      const minimum = readUint32(ctx);
      const value = `${mname} ${rname} ${serial} ${refresh} ${retry} ${expire} ${minimum}`;
      return { expire, minimum, mname, refresh, retry, rname, serial, value };
    }
    case RRTYPE_SRV: {
      const priority = readUint16(ctx);
      const weight = readUint16(ctx);
      const port = readUint16(ctx);
      const target = readName(ctx);
      const value = `${priority} ${weight} ${port} ${target}`;
      return { port, priority, target, value, weight };
    }
    case RRTYPE_TXT: {
      const texts: string[] = [];
      let pos = rdStart;
      while (pos < rdEnd) {
        const txtLen = ctx.data[pos]!;
        pos++;
        if (pos + txtLen > rdEnd) break;
        texts.push(new TextDecoder().decode(ctx.data.subarray(pos, pos + txtLen)));
        pos += txtLen;
      }
      const fullText = texts.join("");
      return { text: fullText, value: fullText };
    }
    default: {
      const raw = Array.from(ctx.data.subarray(rdStart, rdEnd));
      const value = raw.map((b) => b.toString(16).padStart(2, "0")).join(" ");
      return { raw, value };
    }
  }
}

function readAnswers(ctx: ParseContext, count: number): DnsAnswer[] {
  const answers: DnsAnswer[] = [];
  for (let i = 0; i < count; i++) {
    const name = readName(ctx);
    const type = readUint16(ctx);
    const cls = readUint16(ctx);
    const ttl = readUint32(ctx);
    const rdLength = readUint16(ctx);
    const rdStart = ctx.offset;

    const data = parseRdata(ctx, type, rdLength);
    ctx.offset = rdStart + rdLength;

    answers.push({ class: cls, data, name, ttl, type, value: data.value });
  }
  return answers;
}

function readName(ctx: ParseContext): string {
  const parts: string[] = [];
  const visited = new Set<number>();
  let savedOffset: null | number = null;
  let currentOffset = ctx.offset;

  while (true) {
    if (currentOffset >= ctx.data.byteLength) {
      throw new Error("DNS 名称解析越界");
    }
    const len = ctx.data[currentOffset]!;

    if (len === 0) {
      currentOffset++;
      break;
    }

    if ((len & 0xc0) === 0xc0) {
      if (currentOffset + 1 >= ctx.data.byteLength) {
        throw new Error("DNS 压缩指针越界");
      }
      savedOffset ??= currentOffset + 2;
      const ptrOffset = ((len & 0x3f) << 8) | ctx.data[currentOffset + 1]!;
      if (visited.has(ptrOffset)) {
        throw new Error("DNS 压缩指针循环");
      }
      visited.add(ptrOffset);
      currentOffset = ptrOffset;
      continue;
    }

    currentOffset++;
    if (currentOffset + len > ctx.data.byteLength) {
      throw new Error("DNS 标签越界");
    }
    const label = new TextDecoder().decode(ctx.data.subarray(currentOffset, currentOffset + len));
    parts.push(label);
    currentOffset += len;
  }

  ctx.offset = savedOffset ?? currentOffset;
  return parts.join(".");
}

function readUint16(ctx: ParseContext): number {
  const val = ctx.view.getUint16(ctx.offset);
  ctx.offset += 2;
  return val;
}

function readUint32(ctx: ParseContext): number {
  const val = ctx.view.getUint32(ctx.offset);
  ctx.offset += 4;
  return val;
}
