//
//  DNSQuestion.swift
//  punchnet
//
//  Created by 安礼成 on 2026/4/10.
//
import Foundation
import Network

// MARK: - DNS 協議模型
struct DNSQuestion {
    let name: String
    let type: UInt16
    let qclass: UInt16
}

struct DNSResourceRecord {
    let name: String
    let type: UInt16
    let rclass: UInt16
    let ttl: UInt32
    let rdLength: UInt16
    let rdata: Data
}

struct DNSMessage {
    var transactionID: UInt16
    var flags: UInt16
    var questions: [DNSQuestion] = []
    var answers: [DNSResourceRecord] = []
    
    var isResponse: Bool {
        (flags & 0x8000) != 0
    }
}

struct DNSQuerySummary {
    let transactionID: UInt16
    let name: String
    let type: UInt16
    let qclass: UInt16
}

// MARK: - DNS 完整解析器
final class DNSParser {
    private let data: Data
    private var offset: Int = 0

    init(data: Data, offset: Int) {
        self.data = data
        self.offset = offset
    }

    func parse() -> DNSMessage? {
        guard data.count >= 12 + self.offset else {
            return nil
        }
        
        let id = readUInt16()
        let flags = readUInt16()
        let qdCount = readUInt16()
        let anCount = readUInt16()
        let _ = readUInt16() // NSCount
        let _ = readUInt16() // ARCount
        
        var message = DNSMessage(transactionID: id, flags: flags)
        
        for _ in 0..<qdCount {
            if let q = parseQuestion() {
                message.questions.append(q)
            }
        }
        
        for _ in 0..<anCount {
            if let rr = parseRR() {
                message.answers.append(rr)
            }
        }
        
        return message
    }

    private func parseName() -> String {
        var parts: [String] = []
        var jumped = false
        var nextOffset = 0
        var currentOffset = self.offset

        while currentOffset < data.count {
            let length = Int(data[currentOffset])
            if length == 0 {
                currentOffset += 1
                break
            }
            if (length & 0xC0) == 0xC0 {
                let pointer = Int(UInt16(data[currentOffset] & 0x3F) << 8 | UInt16(data[currentOffset + 1]))
                if !jumped {
                    nextOffset = currentOffset + 2
                    jumped = true
                }
                currentOffset = pointer
            } else {
                currentOffset += 1
                if let label = String(data: data.subdata(in: currentOffset..<currentOffset+length), encoding: .ascii) {
                    parts.append(label)
                }
                currentOffset += length
            }
        }
        self.offset = jumped ? nextOffset : currentOffset
        return parts.joined(separator: ".")
    }

    private func parseQuestion() -> DNSQuestion? {
        let name = parseName()
        return DNSQuestion(name: name, type: readUInt16(), qclass: readUInt16())
    }

    private func parseRR() -> DNSResourceRecord? {
        let name = parseName()
        let type = readUInt16()
        let rclass = readUInt16()
        let ttl = readUInt32()
        let rdLength = readUInt16()
        guard offset + Int(rdLength) <= data.count else { return nil }
        let rdata = data.subdata(in: offset..<offset + Int(rdLength))
        offset += Int(rdLength)
        return DNSResourceRecord(name: name, type: type, rclass: rclass, ttl: ttl, rdLength: rdLength, rdata: rdata)
    }

    private func readUInt16() -> UInt16 {
        guard offset + 2 <= data.count else { return 0 }
        let val = UInt16(data[offset]) << 8 | UInt16(data[offset + 1])
        offset += 2
        return val
    }

    private func readUInt32() -> UInt32 {
        guard offset + 4 <= data.count else { return 0 }
        let val = UInt32(data[offset]) << 24 | UInt32(data[offset+1]) << 16 | UInt32(data[offset+2]) << 8 | UInt32(data[offset+3])
        offset += 4
        return val
    }
}

extension DNSParser {
    static func parseFirstQuestion(data: Data, offset: Int) -> DNSQuerySummary? {
        guard offset >= 0, data.count >= offset + 12 else {
            return nil
        }

        return data.withUnsafeBytes { rawBuffer -> DNSQuerySummary? in
            let bytes = rawBuffer.bindMemory(to: UInt8.self)
            guard let baseAddress = bytes.baseAddress else {
                return nil
            }

            func readUInt16(at index: Int) -> UInt16 {
                UInt16(baseAddress[index]) << 8 | UInt16(baseAddress[index + 1])
            }

            let transactionID = readUInt16(at: offset)
            let questionCount = readUInt16(at: offset + 4)
            guard questionCount > 0 else {
                return nil
            }

            var cursor = offset + 12
            guard let name = parseName(
                baseAddress: baseAddress,
                count: data.count,
                messageStart: offset,
                cursor: &cursor
            ), cursor + 4 <= data.count else {
                return nil
            }

            return DNSQuerySummary(
                transactionID: transactionID,
                name: name,
                type: readUInt16(at: cursor),
                qclass: readUInt16(at: cursor + 2)
            )
        }
    }

    private static func parseName(
        baseAddress: UnsafePointer<UInt8>,
        count: Int,
        messageStart: Int,
        cursor: inout Int
    ) -> String? {
        var currentOffset = cursor
        var resumeOffset: Int?
        var jumpCount = 0
        var name = ""

        while currentOffset < count {
            let length = Int(baseAddress[currentOffset])

            if length == 0 {
                currentOffset += 1
                cursor = resumeOffset ?? currentOffset
                return name
            }

            if (length & 0xC0) == 0xC0 {
                guard currentOffset + 1 < count else {
                    return nil
                }

                let pointer = Int(UInt16(baseAddress[currentOffset] & 0x3F) << 8 | UInt16(baseAddress[currentOffset + 1]))
                let targetOffset = messageStart + pointer
                guard targetOffset < count, jumpCount < 8 else {
                    return nil
                }

                if resumeOffset == nil {
                    resumeOffset = currentOffset + 2
                }
                currentOffset = targetOffset
                jumpCount += 1
                continue
            }

            guard (length & 0xC0) == 0, length <= 63 else {
                return nil
            }

            let labelStart = currentOffset + 1
            guard labelStart + length <= count else {
                return nil
            }

            if !name.isEmpty {
                name.append(".")
            }

            let labelBuffer = UnsafeBufferPointer(start: baseAddress.advanced(by: labelStart), count: length)
            name.append(String(decoding: labelBuffer, as: UTF8.self))
            currentOffset = labelStart + length
        }

        return nil
    }
}