const std = @import("std");
const testing = std.testing;

const packet = @import("packet");
const Header = @import("header").Header;
const Name = @import("name").Name;
const Question = @import("question").Question;
const types = @import("types");

// ============================================================================
// Real DNS packet captures for testing
// ============================================================================

/// Standard A query for google.com
/// Captured from: dig google.com A
const GOOGLE_A_QUERY = [_]u8{
    // Header
    0xAB, 0xCD, // ID: 0xABCD
    0x01, 0x00, // Flags: standard query, RD=1
    0x00, 0x01, // QDCOUNT: 1
    0x00, 0x00, // ANCOUNT: 0
    0x00, 0x00, // NSCOUNT: 0
    0x00, 0x00, // ARCOUNT: 0
    // Question: google.com A IN
    0x06, 'g', 'o', 'o', 'g', 'l', 'e',
    0x03, 'c', 'o', 'm',
    0x00, // null terminator
    0x00, 0x01, // QTYPE: A (1)
    0x00, 0x01, // QCLASS: IN (1)
};

/// AAAA query for example.org
const EXAMPLE_AAAA_QUERY = [_]u8{
    // Header
    0x12, 0x34, // ID
    0x01, 0x00, // Flags: standard query, RD=1
    0x00, 0x01, // QDCOUNT: 1
    0x00, 0x00, // ANCOUNT: 0
    0x00, 0x00, // NSCOUNT: 0
    0x00, 0x00, // ARCOUNT: 0
    // Question: example.org AAAA IN
    0x07, 'e', 'x', 'a', 'm', 'p', 'l', 'e',
    0x03, 'o', 'r', 'g',
    0x00,
    0x00, 0x1C, // QTYPE: AAAA (28)
    0x00, 0x01, // QCLASS: IN
};

/// Response with A record
const SIMPLE_A_RESPONSE = [_]u8{
    // Header
    0xAB, 0xCD, // ID
    0x81, 0x80, // Flags: response, RD=1, RA=1
    0x00, 0x01, // QDCOUNT: 1
    0x00, 0x01, // ANCOUNT: 1
    0x00, 0x00, // NSCOUNT: 0
    0x00, 0x00, // ARCOUNT: 0
    // Question: google.com A IN (with compression)
    0x06, 'g', 'o', 'o', 'g', 'l', 'e',
    0x03, 'c', 'o', 'm',
    0x00,
    0x00, 0x01, // QTYPE: A
    0x00, 0x01, // QCLASS: IN
    // Answer: A record using compression pointer
    0xC0, 0x0C, // Name pointer to offset 12 (google.com)
    0x00, 0x01, // TYPE: A
    0x00, 0x01, // CLASS: IN
    0x00, 0x00, 0x01, 0x2C, // TTL: 300 seconds
    0x00, 0x04, // RDLENGTH: 4
    0xD8, 0x3A, 0xD3, 0x8E, // RDATA: 216.58.211.142
};

/// Response with CNAME chain
const CNAME_RESPONSE = [_]u8{
    // Header
    0x55, 0x66, // ID
    0x81, 0x80, // Flags: response, RD=1, RA=1
    0x00, 0x01, // QDCOUNT: 1
    0x00, 0x02, // ANCOUNT: 2 (CNAME + A)
    0x00, 0x00, // NSCOUNT: 0
    0x00, 0x00, // ARCOUNT: 0
    // Question: www.example.com A IN
    0x03, 'w', 'w', 'w',
    0x07, 'e', 'x', 'a', 'm', 'p', 'l', 'e',
    0x03, 'c', 'o', 'm',
    0x00,
    0x00, 0x01, // QTYPE: A
    0x00, 0x01, // QCLASS: IN
    // Answer 1: CNAME www.example.com -> example.com
    0xC0, 0x0C, // Name pointer to www.example.com
    0x00, 0x05, // TYPE: CNAME
    0x00, 0x01, // CLASS: IN
    0x00, 0x00, 0x0E, 0x10, // TTL: 3600
    0x00, 0x02, // RDLENGTH: 2 (compression pointer)
    0xC0, 0x10, // RDATA: pointer to example.com
    // Answer 2: A record for example.com
    0xC0, 0x10, // Name pointer to example.com
    0x00, 0x01, // TYPE: A
    0x00, 0x01, // CLASS: IN
    0x00, 0x00, 0x01, 0x2C, // TTL: 300
    0x00, 0x04, // RDLENGTH: 4
    0x5D, 0xB8, 0xD8, 0x22, // RDATA: 93.184.216.34
};

/// Malformed packet - truncated header
const MALFORMED_TRUNCATED = [_]u8{
    0x12, 0x34, // Only 2 bytes, header needs 12
};

/// Malformed packet - invalid compression pointer (loop)
const MALFORMED_COMPRESSION_LOOP = [_]u8{
    // Header
    0x00, 0x01,
    0x01, 0x00,
    0x00, 0x01,
    0x00, 0x00,
    0x00, 0x00,
    0x00, 0x00,
    // Question with self-referencing pointer
    0xC0, 0x0C, // Points to itself
    0x00, 0x01,
    0x00, 0x01,
};

// ============================================================================
// Header Tests
// ============================================================================

test "Header.parse - standard query" {
    const h = try Header.parse(&GOOGLE_A_QUERY);

    try testing.expectEqual(@as(u16, 0xABCD), h.id);
    try testing.expect(!h.qr); // Query, not response
    try testing.expectEqual(types.OpCode.Query, h.opcode);
    try testing.expect(!h.aa); // Not authoritative
    try testing.expect(!h.tc); // Not truncated
    try testing.expect(h.rd); // Recursion desired
    try testing.expect(!h.ra); // Recursion not available (query)
    try testing.expectEqual(types.RCode.NoError, h.rcode);
    try testing.expectEqual(@as(u16, 1), h.qdcount);
    try testing.expectEqual(@as(u16, 0), h.ancount);
    try testing.expectEqual(@as(u16, 0), h.nscount);
    try testing.expectEqual(@as(u16, 0), h.arcount);
}

test "Header.parse - standard response" {
    const h = try Header.parse(&SIMPLE_A_RESPONSE);

    try testing.expectEqual(@as(u16, 0xABCD), h.id);
    try testing.expect(h.qr); // Response
    try testing.expect(h.rd); // RD copied from query
    try testing.expect(h.ra); // Recursion available
    try testing.expectEqual(@as(u16, 1), h.qdcount);
    try testing.expectEqual(@as(u16, 1), h.ancount);
}

test "Header.parse - buffer too small" {
    const result = Header.parse(&MALFORMED_TRUNCATED);
    try testing.expectError(error.BufferTooSmall, result);
}

test "Header.encode - roundtrip" {
    const original = try Header.parse(&GOOGLE_A_QUERY);

    var buf: [12]u8 = undefined;
    original.encode(&buf);

    const decoded = try Header.parse(&buf);
    try testing.expectEqual(original.id, decoded.id);
    try testing.expectEqual(original.qr, decoded.qr);
    try testing.expectEqual(original.opcode, decoded.opcode);
    try testing.expectEqual(original.rd, decoded.rd);
    try testing.expectEqual(original.qdcount, decoded.qdcount);
}

// ============================================================================
// Name Tests
// ============================================================================

test "Name.parse - simple domain" {
    const allocator = testing.allocator;
    const buffer = [_]u8{ 0x06, 'g', 'o', 'o', 'g', 'l', 'e', 0x03, 'c', 'o', 'm', 0x00 };

    const result = try Name.parse(&buffer, &buffer, allocator);
    defer result.name.deinit();

    try testing.expectEqual(@as(usize, 12), result.bytes_read);

    var str_buf: [256]u8 = undefined;
    const str = result.name.toStringBuf(&str_buf).?;
    try testing.expectEqualStrings("google.com", str);
}

test "Name.parse - root domain" {
    const allocator = testing.allocator;
    const buffer = [_]u8{0x00}; // Just null byte = root

    const result = try Name.parse(&buffer, &buffer, allocator);
    defer result.name.deinit();

    try testing.expectEqual(@as(usize, 1), result.bytes_read);

    var str_buf: [256]u8 = undefined;
    const str = result.name.toStringBuf(&str_buf).?;
    try testing.expectEqualStrings("", str);
}

test "Name.parse - compression pointer" {
    const allocator = testing.allocator;

    // Full packet with question followed by answer using compression
    const result = try Name.parse(SIMPLE_A_RESPONSE[27..], &SIMPLE_A_RESPONSE, allocator);
    defer result.name.deinit();

    var str_buf: [256]u8 = undefined;
    const str = result.name.toStringBuf(&str_buf).?;
    try testing.expectEqualStrings("google.com", str);
}

test "Name.parse - max label length (63)" {
    const allocator = testing.allocator;
    var buffer: [67]u8 = undefined;
    buffer[0] = 63; // Label length = max
    for (1..64) |i| {
        buffer[i] = 'a';
    }
    buffer[64] = 0x03;
    buffer[65] = 'c';
    buffer[66] = 'o';
    // Would need more bytes for full domain, but testing max label

    // This should work (63 is max label length)
    const result = Name.parse(buffer[0..67], buffer[0..67], allocator);
    if (result) |r| {
        r.name.deinit();
    } else |_| {
        // May fail due to incomplete buffer, which is fine
    }
}

test "Name.parse - label too long (64+)" {
    const allocator = testing.allocator;
    var buffer: [68]u8 = undefined;
    buffer[0] = 64; // Label length > 63 is invalid (and not a pointer)
    @memset(buffer[1..65], 'a');
    buffer[65] = 0x00;

    const result = Name.parse(&buffer, &buffer, allocator);
    try testing.expectError(error.InvalidLabel, result);
}

// ============================================================================
// Question Tests
// ============================================================================

test "Question.parse - A record query" {
    const allocator = testing.allocator;

    const result = try Question.parse(GOOGLE_A_QUERY[12..], &GOOGLE_A_QUERY, allocator);
    defer result.question.deinit();

    try testing.expectEqual(types.QType.A, result.question.qtype);
    try testing.expectEqual(types.QClass.IN, result.question.qclass);

    var str_buf: [256]u8 = undefined;
    const domain = result.question.name.toStringBuf(&str_buf).?;
    try testing.expectEqualStrings("google.com", domain);
}

test "Question.parse - AAAA record query" {
    const allocator = testing.allocator;

    const result = try Question.parse(EXAMPLE_AAAA_QUERY[12..], &EXAMPLE_AAAA_QUERY, allocator);
    defer result.question.deinit();

    try testing.expectEqual(types.QType.AAAA, result.question.qtype);
    try testing.expectEqual(types.QClass.IN, result.question.qclass);

    var str_buf: [256]u8 = undefined;
    const domain = result.question.name.toStringBuf(&str_buf).?;
    try testing.expectEqualStrings("example.org", domain);
}

// ============================================================================
// Packet Tests
// ============================================================================

test "Packet.parse - simple query" {
    const allocator = testing.allocator;

    var pkt = try packet.Packet.parse(&GOOGLE_A_QUERY, allocator);
    defer pkt.deinit();

    try testing.expectEqual(@as(u16, 0xABCD), pkt.header.id);
    try testing.expect(!pkt.header.qr);
    try testing.expectEqual(@as(usize, 1), pkt.questions.len);
    try testing.expectEqual(@as(usize, 0), pkt.answers.len);

    var str_buf: [256]u8 = undefined;
    const domain = pkt.questions[0].name.toStringBuf(&str_buf).?;
    try testing.expectEqualStrings("google.com", domain);
    try testing.expectEqual(types.QType.A, pkt.questions[0].qtype);
}

test "Packet.parse - response with A record" {
    const allocator = testing.allocator;

    var pkt = try packet.Packet.parse(&SIMPLE_A_RESPONSE, allocator);
    defer pkt.deinit();

    try testing.expect(pkt.header.qr);
    try testing.expectEqual(@as(usize, 1), pkt.questions.len);
    try testing.expectEqual(@as(usize, 1), pkt.answers.len);

    // Check the A record
    const answer = pkt.answers[0];
    try testing.expectEqual(types.QType.A, answer.rtype);
    try testing.expectEqual(@as(u32, 300), answer.ttl);

    // Verify IP address
    const ip = answer.getA().?;
    try testing.expectEqual([4]u8{ 0xD8, 0x3A, 0xD3, 0x8E }, ip);
}

test "Packet.parse - CNAME chain" {
    const allocator = testing.allocator;

    var pkt = try packet.Packet.parse(&CNAME_RESPONSE, allocator);
    defer pkt.deinit();

    try testing.expectEqual(@as(usize, 2), pkt.answers.len);

    // First answer should be CNAME
    try testing.expectEqual(types.QType.CNAME, pkt.answers[0].rtype);

    // Second answer should be A record
    try testing.expectEqual(types.QType.A, pkt.answers[1].rtype);

    const ip = pkt.answers[1].getA().?;
    try testing.expectEqual([4]u8{ 0x5D, 0xB8, 0xD8, 0x22 }, ip);
}

test "Packet.parse - truncated packet" {
    const allocator = testing.allocator;

    const result = packet.Packet.parse(&MALFORMED_TRUNCATED, allocator);
    try testing.expectError(error.HeaderParseError, result);
}

test "Packet.encode - roundtrip" {
    const allocator = testing.allocator;

    var original = try packet.Packet.parse(&GOOGLE_A_QUERY, allocator);
    defer original.deinit();

    var buf: [512]u8 = undefined;
    const encoded_len = try original.encode(&buf);

    var decoded = try packet.Packet.parse(buf[0..encoded_len], allocator);
    defer decoded.deinit();

    try testing.expectEqual(original.header.id, decoded.header.id);
    try testing.expectEqual(original.questions.len, decoded.questions.len);
}

test "Packet.createBlockedResponse - creates valid response" {
    const allocator = testing.allocator;

    var query = try packet.Packet.parse(&GOOGLE_A_QUERY, allocator);
    defer query.deinit();

    var response = try packet.Packet.createBlockedResponse(&query, allocator);
    defer response.deinit();

    try testing.expect(response.header.qr); // Is response
    try testing.expectEqual(query.header.id, response.header.id);
    try testing.expectEqual(@as(usize, 1), response.answers.len);

    // Should return 0.0.0.0 for blocked
    const ip = response.answers[0].getA().?;
    try testing.expectEqual([4]u8{ 0, 0, 0, 0 }, ip);
}

test "Packet.createNxdomainResponse - creates valid NXDOMAIN" {
    const allocator = testing.allocator;

    var query = try packet.Packet.parse(&GOOGLE_A_QUERY, allocator);
    defer query.deinit();

    var response = try packet.Packet.createNxdomainResponse(&query, allocator);
    defer response.deinit();

    try testing.expect(response.header.qr);
    try testing.expectEqual(types.RCode.NXDomain, response.header.rcode);
    try testing.expectEqual(@as(usize, 0), response.answers.len);
}

// ============================================================================
// Edge Cases
// ============================================================================

test "DNS max name length (253 chars)" {
    const allocator = testing.allocator;

    // Build a name with max length: 63.63.63.63 = 253 chars + labels
    var buffer: [512]u8 = undefined;
    var pos: usize = 12; // Skip header

    // Add 4 labels of 63 chars each
    for (0..4) |_| {
        buffer[pos] = 63;
        pos += 1;
        @memset(buffer[pos .. pos + 63], 'a');
        pos += 63;
    }
    buffer[pos] = 0; // Null terminator
    pos += 1;

    // Add qtype and qclass
    buffer[pos] = 0x00;
    buffer[pos + 1] = 0x01;
    buffer[pos + 2] = 0x00;
    buffer[pos + 3] = 0x01;
    pos += 4;

    // Set up header
    @memset(buffer[0..12], 0);
    buffer[2] = 0x01; // RD=1
    buffer[5] = 0x01; // QDCOUNT=1

    const pkt_result = packet.Packet.parse(buffer[0..pos], allocator);
    // This may fail due to name being too long (4*63=252 + dots)
    if (pkt_result) |*pkt| {
        pkt.deinit();
    } else |_| {}
}

test "DNS various query types" {
    const allocator = testing.allocator;
    const query_types = [_]types.QType{ .A, .AAAA, .CNAME, .MX, .NS, .TXT, .SOA, .PTR, .SRV };

    for (query_types) |qtype| {
        var buf: [64]u8 = undefined;
        @memset(buf[0..12], 0);
        buf[2] = 0x01; // RD
        buf[5] = 0x01; // QDCOUNT

        // Simple question: a.b
        buf[12] = 0x01;
        buf[13] = 'a';
        buf[14] = 0x01;
        buf[15] = 'b';
        buf[16] = 0x00;

        const qtype_val = @intFromEnum(qtype);
        buf[17] = @intCast((qtype_val >> 8) & 0xFF);
        buf[18] = @intCast(qtype_val & 0xFF);
        buf[19] = 0x00;
        buf[20] = 0x01; // CLASS IN

        var pkt = try packet.Packet.parse(buf[0..21], allocator);
        defer pkt.deinit();

        try testing.expectEqual(qtype, pkt.questions[0].qtype);
    }
}