import { describe, expect, it } from "bun:test";

import {
  buildQuery,
  parseResponse,
  rcodeName,
  rrtTypeByName,
  rrtTypeName,
} from "../../../../../src/server/checker/runner/dns/codec";

function buildName(name: string): Uint8Array {
  const parts: number[] = [];
  for (const label of name.split(".")) {
    const encoded = new TextEncoder().encode(label);
    parts.push(encoded.length);
    parts.push(...encoded);
  }
  parts.push(0);
  return new Uint8Array(parts);
}

function buildResponse(options: {
  answers?: Array<{ class?: number; name: string; rdata: Uint8Array; ttl: number; type: number }>;
  flags?: { aa?: boolean; ad?: boolean; ra?: boolean; rd?: boolean; tc?: boolean };
  questions?: Array<{ name: string; qclass?: number; qtype: number }>;
  rcode?: number;
}): Uint8Array {
  const questions = options.questions ?? [];
  const answers = options.answers ?? [];
  const id = 0x1234;
  let flags = 0x8000;
  if (options.flags?.aa) flags |= 0x0400;
  if (options.flags?.tc) flags |= 0x0200;
  if (options.flags?.rd) flags |= 0x0100;
  if (options.flags?.ra) flags |= 0x0080;
  if (options.flags?.ad) flags |= 0x0020;
  flags |= (options.rcode ?? 0) & 0x000f;

  const header = new Uint8Array(12);
  const hv = new DataView(header.buffer);
  hv.setUint16(0, id);
  hv.setUint16(2, flags);
  hv.setUint16(4, questions.length);
  hv.setUint16(6, answers.length);
  hv.setUint16(8, 0);
  hv.setUint16(10, 0);

  const qParts: Uint8Array[] = [];
  for (const q of questions) {
    const nameBytes = buildName(q.name);
    const qtype = new Uint8Array(4);
    const qv = new DataView(qtype.buffer);
    qv.setUint16(0, q.qtype);
    qv.setUint16(2, q.qclass ?? 1);
    qParts.push(nameBytes, qtype);
  }

  const aParts: Uint8Array[] = [];
  for (const a of answers) {
    const nameBytes = buildName(a.name);
    const rrHead = new Uint8Array(10);
    const rv = new DataView(rrHead.buffer);
    rv.setUint16(0, a.type);
    rv.setUint16(2, a.class ?? 1);
    rv.setUint32(4, a.ttl);
    rv.setUint16(8, a.rdata.length);
    aParts.push(nameBytes, rrHead, a.rdata);
  }

  const allParts = [header, ...qParts, ...aParts];
  const totalLen = allParts.reduce((s, p) => s + p.length, 0);
  const result = new Uint8Array(totalLen);
  let offset = 0;
  for (const part of allParts) {
    result.set(part, offset);
    offset += part.length;
  }
  return result;
}

describe("buildQuery", () => {
  it("produces a query buffer with correct header (QR=0, OPCODE=0, RD flag)", () => {
    const buf = buildQuery("example.com", 1, true);
    const view = new DataView(buf.buffer, buf.byteOffset, buf.byteLength);
    expect(buf.length).toBeGreaterThan(12);
    const flags = view.getUint16(2);
    expect(flags & 0x8000).toBe(0);
    expect(flags & 0x7800).toBe(0);
    expect(flags & 0x0100).toBe(0x0100);
    expect(view.getUint16(4)).toBe(1);
    expect(view.getUint16(6)).toBe(0);
    expect(view.getUint16(8)).toBe(0);
    expect(view.getUint16(10)).toBe(0);
  });

  it("encodes domain name correctly (labels with length prefixes, null terminator)", () => {
    const buf = buildQuery("example.com", 1, false);
    const nameBytes = buf.subarray(12, buf.length - 4);
    const expected = buildName("example.com");
    expect(Array.from(nameBytes)).toEqual(Array.from(expected));
  });

  it("encodes trailing-dot FQDN without an extra root label", () => {
    const normal = buildQuery("example.com", 1, false);
    const fqdn = buildQuery("example.com.", 1, false);
    expect(Array.from(fqdn.subarray(12))).toEqual(Array.from(normal.subarray(12)));
  });

  it("encodes root domain", () => {
    const buf = buildQuery(".", 1, false);
    expect(Array.from(buf.subarray(12, buf.length - 4))).toEqual([0]);
  });

  it("rejects non-ASCII names instead of writing malformed UTF-8 labels", () => {
    expect(() => buildQuery("é.example", 1, false)).toThrow("ASCII/Punycode");
  });

  it("rejects labels longer than 63 bytes", () => {
    expect(() => buildQuery(`${"a".repeat(64)}.example`, 1, false)).toThrow("63 字节");
  });

  it("sets correct QTYPE and QCLASS", () => {
    const buf = buildQuery("example.com", 28, false);
    const view = new DataView(buf.buffer, buf.byteOffset, buf.byteLength);
    const qtypeOffset = buf.length - 4;
    expect(view.getUint16(qtypeOffset)).toBe(28);
    expect(view.getUint16(qtypeOffset + 2)).toBe(1);
  });

  it("sets RD=1 when recursionDesired=true, RD=0 when false", () => {
    const bufRD = buildQuery("example.com", 1, true);
    const viewRD = new DataView(bufRD.buffer, bufRD.byteOffset, bufRD.byteLength);
    expect(viewRD.getUint16(2) & 0x0100).toBe(0x0100);

    const bufNoRD = buildQuery("example.com", 1, false);
    const viewNoRD = new DataView(bufNoRD.buffer, bufNoRD.byteOffset, bufNoRD.byteLength);
    expect(viewNoRD.getUint16(2) & 0x0100).toBe(0);
  });
});

describe("parseResponse", () => {
  it("parses a minimal NOERROR response (no answers)", () => {
    const resp = buildResponse({
      questions: [{ name: "example.com", qtype: 1 }],
      rcode: 0,
    });
    const result = parseResponse(resp);
    expect(result.header.rcode).toBe(0);
    expect(result.header.answerCount).toBe(0);
    expect(result.answers).toHaveLength(0);
    expect(result.questions).toHaveLength(1);
    expect(result.questions[0]!.name).toBe("example.com");
    expect(result.questions[0]!.qtype).toBe(1);
  });

  it("parses response with single A record", () => {
    const rdata = new Uint8Array([93, 184, 216, 34]);
    const resp = buildResponse({
      answers: [{ name: "example.com", rdata, ttl: 300, type: 1 }],
      questions: [{ name: "example.com", qtype: 1 }],
      rcode: 0,
    });
    const result = parseResponse(resp);
    expect(result.answers).toHaveLength(1);
    const a = result.answers[0]!;
    expect(a.type).toBe(1);
    expect(a.value).toBe("93.184.216.34");
    expect(a.ttl).toBe(300);
    expect(a.name).toBe("example.com");
  });

  it("parses response with multiple A records", () => {
    const rdata1 = new Uint8Array([1, 1, 1, 1]);
    const rdata2 = new Uint8Array([8, 8, 8, 8]);
    const resp = buildResponse({
      answers: [
        { name: "example.com", rdata: rdata1, ttl: 60, type: 1 },
        { name: "example.com", rdata: rdata2, ttl: 60, type: 1 },
      ],
      rcode: 0,
    });
    const result = parseResponse(resp);
    expect(result.answers).toHaveLength(2);
    expect(result.answers[0]!.value).toBe("1.1.1.1");
    expect(result.answers[1]!.value).toBe("8.8.8.8");
  });

  it("parses AAAA record", () => {
    const rdata = new Uint8Array([
      0x26, 0x07, 0xf8, 0xb0, 0x40, 0x05, 0x08, 0x0c, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
    ]);
    const resp = buildResponse({
      answers: [{ name: "example.com", rdata, ttl: 300, type: 28 }],
    });
    const result = parseResponse(resp);
    expect(result.answers[0]!.value).toBe("2607:f8b0:4005:80c::1");
  });

  it("parses CNAME record", () => {
    const rdata = buildName("www.example.com");
    const resp = buildResponse({
      answers: [{ name: "alias.example.com", rdata, ttl: 300, type: 5 }],
    });
    const result = parseResponse(resp);
    expect(result.answers[0]!.value).toBe("www.example.com");
    expect(result.answers[0]!.data["target"]).toBe("www.example.com");
  });

  it("parses MX record (preference + exchange)", () => {
    const exchange = buildName("mail.example.com");
    const rdata = new Uint8Array(2 + exchange.length);
    rdata[0] = 0x00;
    rdata[1] = 0x0a;
    rdata.set(exchange, 2);
    const resp = buildResponse({
      answers: [{ name: "example.com", rdata, ttl: 3600, type: 15 }],
    });
    const result = parseResponse(resp);
    const mx = result.answers[0]!;
    expect(mx.value).toBe("10 mail.example.com");
    expect(mx.data["preference"]).toBe(10);
    expect(mx.data["exchange"]).toBe("mail.example.com");
  });

  it("parses TXT record with single character-string", () => {
    const txt = new TextEncoder().encode("v=spf1 include:_spf.example.com ~all");
    const rdata = new Uint8Array(1 + txt.length);
    rdata[0] = txt.length;
    rdata.set(txt, 1);
    const resp = buildResponse({
      answers: [{ name: "example.com", rdata, ttl: 300, type: 16 }],
    });
    const result = parseResponse(resp);
    expect(result.answers[0]!.value).toBe("v=spf1 include:_spf.example.com ~all");
  });

  it("parses TXT record with multiple character-strings", () => {
    const s1 = new TextEncoder().encode("hello ");
    const s2 = new TextEncoder().encode("world");
    const rdata = new Uint8Array(1 + s1.length + 1 + s2.length);
    rdata[0] = s1.length;
    rdata.set(s1, 1);
    rdata[1 + s1.length] = s2.length;
    rdata.set(s2, 1 + s1.length + 1);
    const resp = buildResponse({
      answers: [{ name: "example.com", rdata, ttl: 300, type: 16 }],
    });
    const result = parseResponse(resp);
    expect(result.answers[0]!.value).toBe("hello world");
  });

  it("parses NS record", () => {
    const rdata = buildName("ns1.example.com");
    const resp = buildResponse({
      answers: [{ name: "example.com", rdata, ttl: 86400, type: 2 }],
    });
    const result = parseResponse(resp);
    expect(result.answers[0]!.value).toBe("ns1.example.com");
    expect(result.answers[0]!.data["nsdname"]).toBe("ns1.example.com");
  });

  it("parses SOA record (all 7 fields)", () => {
    const mname = buildName("ns1.example.com");
    const rname = buildName("admin.example.com");
    const fixed = new Uint8Array(20);
    const fv = new DataView(fixed.buffer);
    fv.setUint32(0, 2024010101);
    fv.setUint32(4, 3600);
    fv.setUint32(8, 900);
    fv.setUint32(12, 604800);
    fv.setUint32(16, 86400);
    const rdata = new Uint8Array(mname.length + rname.length + fixed.length);
    rdata.set(mname, 0);
    rdata.set(rname, mname.length);
    rdata.set(fixed, mname.length + rname.length);
    const resp = buildResponse({
      answers: [{ name: "example.com", rdata, ttl: 3600, type: 6 }],
    });
    const result = parseResponse(resp);
    const soa = result.answers[0]!;
    expect(soa.data["mname"]).toBe("ns1.example.com");
    expect(soa.data["rname"]).toBe("admin.example.com");
    expect(soa.data["serial"]).toBe(2024010101);
    expect(soa.data["refresh"]).toBe(3600);
    expect(soa.data["retry"]).toBe(900);
    expect(soa.data["expire"]).toBe(604800);
    expect(soa.data["minimum"]).toBe(86400);
    expect(soa.value).toBe("ns1.example.com admin.example.com 2024010101 3600 900 604800 86400");
  });

  it("parses SRV record (priority, weight, port, target)", () => {
    const target = buildName("server.example.com");
    const rdata = new Uint8Array(6 + target.length);
    const rv = new DataView(rdata.buffer);
    rv.setUint16(0, 10);
    rv.setUint16(2, 20);
    rv.setUint16(4, 443);
    rdata.set(target, 6);
    const resp = buildResponse({
      answers: [{ name: "_https._tcp.example.com", rdata, ttl: 300, type: 33 }],
    });
    const result = parseResponse(resp);
    const srv = result.answers[0]!;
    expect(srv.data["priority"]).toBe(10);
    expect(srv.data["weight"]).toBe(20);
    expect(srv.data["port"]).toBe(443);
    expect(srv.data["target"]).toBe("server.example.com");
    expect(srv.value).toBe("10 20 443 server.example.com");
  });

  it("parses PTR record", () => {
    const rdata = buildName("host.example.com");
    const resp = buildResponse({
      answers: [{ name: "1.0.0.127.in-addr.arpa", rdata, ttl: 300, type: 12 }],
      questions: [{ name: "1.0.0.127.in-addr.arpa", qtype: 12 }],
    });
    const result = parseResponse(resp);
    expect(result.answers[0]!.value).toBe("host.example.com");
    expect(result.answers[0]!.data["ptrdname"]).toBe("host.example.com");
  });

  it("parses CAA record (flags, tag, value)", () => {
    const tag = new TextEncoder().encode("issue");
    const value = new TextEncoder().encode("letsencrypt.org");
    const rdata = new Uint8Array(2 + tag.length + value.length);
    rdata[0] = 0;
    rdata[1] = tag.length;
    rdata.set(tag, 2);
    rdata.set(value, 2 + tag.length);
    const resp = buildResponse({
      answers: [{ name: "example.com", rdata, ttl: 300, type: 257 }],
    });
    const result = parseResponse(resp);
    const caa = result.answers[0]!;
    expect(caa.data["flags"]).toBe(0);
    expect(caa.data["tag"]).toBe("issue");
    expect(caa.data["valueStr"]).toBe("letsencrypt.org");
    expect(caa.value).toBe("0 issue letsencrypt.org");
  });

  it("parses response with DNS name compression (pointer)", () => {
    const nameBytes = buildName("example.com");
    const nameLen = nameBytes.length;
    const headerSize = 12;
    const qSectionEnd = headerSize + nameLen + 4;

    const totalSize = qSectionEnd + nameLen + 10 + 4;
    const buf = new Uint8Array(totalSize);
    const view = new DataView(buf.buffer);

    view.setUint16(0, 0xabcd);
    view.setUint16(2, 0x8180);
    view.setUint16(4, 1);
    view.setUint16(6, 1);
    view.setUint16(8, 0);
    view.setUint16(10, 0);

    buf.set(nameBytes, headerSize);
    view.setUint16(headerSize + nameLen, 1);
    view.setUint16(headerSize + nameLen + 2, 1);

    const ptrByte1 = 0xc0 | ((headerSize >> 8) & 0x3f);
    const ptrByte2 = headerSize & 0xff;
    buf[qSectionEnd] = ptrByte1;
    buf[qSectionEnd + 1] = ptrByte2;
    view.setUint16(qSectionEnd + 2, 1);
    view.setUint16(qSectionEnd + 4, 1);
    view.setUint32(qSectionEnd + 6, 300);
    view.setUint16(qSectionEnd + 10, 4);
    buf[qSectionEnd + 12] = 93;
    buf[qSectionEnd + 13] = 184;
    buf[qSectionEnd + 14] = 216;
    buf[qSectionEnd + 15] = 34;

    const result = parseResponse(buf);
    expect(result.answers[0]!.name).toBe("example.com");
    expect(result.answers[0]!.value).toBe("93.184.216.34");
  });

  it("correctly reports flags (AA, RA, RD, TC, AD)", () => {
    const resp = buildResponse({
      flags: { aa: true, ad: true, ra: true, rd: true, tc: false },
    });
    const result = parseResponse(resp);
    expect(result.header.flags.authoritative).toBe(true);
    expect(result.header.flags.recursionAvailable).toBe(true);
    expect(result.header.flags.recursionDesired).toBe(true);
    expect(result.header.flags.truncated).toBe(false);
    expect(result.header.flags.authenticatedData).toBe(true);
  });

  it("reports correct rcode", () => {
    const resp = buildResponse({ rcode: 3 });
    const result = parseResponse(resp);
    expect(result.header.rcode).toBe(3);
  });

  it("reports correct answerCount, authorityCount, additionalCount", () => {
    const a1 = new Uint8Array([1, 2, 3, 4]);
    const a2 = new Uint8Array([5, 6, 7, 8]);
    const header = new Uint8Array(12);
    const hv = new DataView(header.buffer);
    hv.setUint16(0, 0x0001);
    hv.setUint16(2, 0x8180);
    hv.setUint16(4, 0);
    hv.setUint16(6, 1);
    hv.setUint16(8, 1);
    hv.setUint16(10, 1);

    const nameBytes = buildName("example.com");
    const buildRR = (rdata: Uint8Array) => {
      const rr = new Uint8Array(nameBytes.length + 10 + rdata.length);
      rr.set(nameBytes, 0);
      const rv = new DataView(rr.buffer);
      rv.setUint16(nameBytes.length, 1);
      rv.setUint16(nameBytes.length + 2, 1);
      rv.setUint32(nameBytes.length + 4, 60);
      rv.setUint16(nameBytes.length + 8, rdata.length);
      rr.set(rdata, nameBytes.length + 10);
      return rr;
    };
    const rr1 = buildRR(a1);
    const rr2 = buildRR(a2);
    const rr3 = buildRR(new Uint8Array([9, 10, 11, 12]));

    const buf = new Uint8Array(12 + rr1.length + rr2.length + rr3.length);
    buf.set(header, 0);
    buf.set(rr1, 12);
    buf.set(rr2, 12 + rr1.length);
    buf.set(rr3, 12 + rr1.length + rr2.length);

    const result = parseResponse(buf);
    expect(result.header.answerCount).toBe(1);
    expect(result.header.authorityCount).toBe(1);
    expect(result.header.additionalCount).toBe(1);
    expect(result.answers).toHaveLength(1);
    expect(result.authorities).toHaveLength(1);
    expect(result.additional).toHaveLength(1);
  });

  it("throws on response shorter than 12 bytes", () => {
    expect(() => parseResponse(new Uint8Array(11))).toThrow();
    expect(() => parseResponse(new Uint8Array(0))).toThrow();
  });

  it("handles unknown record types (raw hex in value)", () => {
    const rdata = new Uint8Array([0xde, 0xad, 0xbe, 0xef]);
    const resp = buildResponse({
      answers: [{ name: "example.com", rdata, ttl: 300, type: 999 }],
    });
    const result = parseResponse(resp);
    expect(result.answers[0]!.value).toBe("de ad be ef");
    expect(result.answers[0]!.data["raw"]).toEqual([0xde, 0xad, 0xbe, 0xef]);
  });
});

describe("rcodeName", () => {
  it("maps known rcode numbers to names", () => {
    expect(rcodeName(0)).toBe("NOERROR");
    expect(rcodeName(1)).toBe("FORMERR");
    expect(rcodeName(2)).toBe("SERVFAIL");
    expect(rcodeName(3)).toBe("NXDOMAIN");
    expect(rcodeName(4)).toBe("NOTIMP");
    expect(rcodeName(5)).toBe("REFUSED");
  });

  it("returns UNKNOWN(n) for unknown codes", () => {
    expect(rcodeName(99)).toBe("UNKNOWN(99)");
    expect(rcodeName(255)).toBe("UNKNOWN(255)");
  });
});

describe("rrtTypeName", () => {
  it("maps known type numbers to names", () => {
    expect(rrtTypeName(1)).toBe("A");
    expect(rrtTypeName(28)).toBe("AAAA");
    expect(rrtTypeName(5)).toBe("CNAME");
    expect(rrtTypeName(15)).toBe("MX");
    expect(rrtTypeName(16)).toBe("TXT");
    expect(rrtTypeName(2)).toBe("NS");
    expect(rrtTypeName(6)).toBe("SOA");
    expect(rrtTypeName(33)).toBe("SRV");
    expect(rrtTypeName(12)).toBe("PTR");
    expect(rrtTypeName(257)).toBe("CAA");
  });
});

describe("rrtTypeByName", () => {
  it("maps known type names to numbers", () => {
    expect(rrtTypeByName("A")).toBe(1);
    expect(rrtTypeByName("AAAA")).toBe(28);
    expect(rrtTypeByName("CNAME")).toBe(5);
    expect(rrtTypeByName("MX")).toBe(15);
    expect(rrtTypeByName("TXT")).toBe(16);
    expect(rrtTypeByName("NS")).toBe(2);
    expect(rrtTypeByName("SOA")).toBe(6);
    expect(rrtTypeByName("SRV")).toBe(33);
    expect(rrtTypeByName("PTR")).toBe(12);
    expect(rrtTypeByName("CAA")).toBe(257);
  });

  it("throws for unknown type name", () => {
    expect(() => rrtTypeByName("UNKNOWN_TYPE")).toThrow();
  });
});
