//
//  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
    }
}

// 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
    }
}