import Foundation

public enum MeshDNSCodec {
    public static func response(to query: Data, aliases: [String: MeshAddress], ttl: UInt32 = 30) throws -> Data {
        guard query.count >= 12, read16(query, 4) == 1 else { throw UMNError.invalidFrame }
        var offset = 12; var visited = Set<Int>()
        let labels = try decodeName(query, offset: &offset, visited: &visited)
        guard offset + 4 <= query.count else { throw UMNError.invalidFrame }
        let type = read16(query, offset), qclass = read16(query, offset + 2); offset += 4
        guard qclass == 1 else { throw UMNError.invalidFrame }
        let question = query.subdata(in: 12..<offset)
        let fqdn = labels.joined(separator: ".").lowercased()
        let aliasName = fqdn.hasSuffix(".mesh") ? String(fqdn.dropLast(5)) : ""
        let address = aliases[aliasName]
        let answer = address != nil && type == 28
        let rcode: UInt16 = address == nil ? 3 : 0
        var output = Data()
        append16(read16(query, 0), to: &output)
        append16(0x8400 | rcode, to: &output) // response + authoritative
        append16(1, to: &output); append16(answer ? 1 : 0, to: &output)
        append16(0, to: &output); append16(0, to: &output); output.append(question)
        if answer, let address {
            append16(0xc00c, to: &output); append16(28, to: &output); append16(1, to: &output)
            append32(ttl, to: &output); append16(16, to: &output); output.append(address.bytes)
        }
        guard output.count <= 512 else { throw UMNError.invalidFrame }
        return output
    }

    private static func decodeName(_ data: Data, offset: inout Int, visited: inout Set<Int>, depth: Int = 0) throws -> [String] {
        guard depth < 16 else { throw UMNError.invalidFrame }
        var labels: [String] = []
        while true {
            guard offset < data.count else { throw UMNError.invalidFrame }
            let length = Int(data[offset]); offset += 1
            if length == 0 { return labels }
            if length & 0xc0 == 0xc0 {
                guard offset < data.count else { throw UMNError.invalidFrame }
                let pointer = ((length & 0x3f) << 8) | Int(data[offset]); offset += 1
                guard pointer >= 12, pointer < data.count, visited.insert(pointer).inserted else { throw UMNError.invalidFrame }
                var nested = pointer
                labels += try decodeName(data, offset: &nested, visited: &visited, depth: depth + 1)
                return labels
            }
            guard length <= 63, offset + length <= data.count,
                  let label = String(data: data.subdata(in: offset..<(offset + length)), encoding: .utf8),
                  !label.isEmpty else { throw UMNError.invalidFrame }
            labels.append(label); offset += length
        }
    }
    private static func read16(_ data: Data, _ offset: Int) -> UInt16 { UInt16(data[offset]) << 8 | UInt16(data[offset + 1]) }
    private static func append16(_ value: UInt16, to data: inout Data) { data.append(UInt8(value >> 8)); data.append(UInt8(value & 255)) }
    private static func append32(_ value: UInt32, to data: inout Data) {
        data.append(UInt8(value >> 24)); data.append(UInt8((value >> 16) & 255)); data.append(UInt8((value >> 8) & 255)); data.append(UInt8(value & 255))
    }
}