import Foundation
import Darwin
import UltraMeshCore

final class MeshDNSServer {
    static let port: UInt16 = 53535
    private let aliases: () -> [String: MeshAddress]
    private var udp: Int32 = -1, tcp: Int32 = -1
    private var running = false

    init(aliases: @escaping () -> [String: MeshAddress]) { self.aliases = aliases }

    func start() throws {
        udp = try bindSocket(type: SOCK_DGRAM); tcp = try bindSocket(type: SOCK_STREAM)
        guard Darwin.listen(tcp, 32) == 0 else { throw POSIXError(.init(rawValue: errno) ?? .EIO) }
        running = true
        DispatchQueue.global(qos: .utility).async { [weak self] in self?.udpLoop() }
        DispatchQueue.global(qos: .utility).async { [weak self] in self?.tcpLoop() }
    }
    func stop() {
        running = false
        if udp >= 0 { Darwin.close(udp); udp = -1 }
        if tcp >= 0 { Darwin.close(tcp); tcp = -1 }
    }

    private func bindSocket(type: Int32) throws -> Int32 {
        let fd = socket(AF_INET, type, 0); guard fd >= 0 else { throw POSIXError(.init(rawValue: errno) ?? .EIO) }
        var yes: Int32 = 1; setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &yes, socklen_t(MemoryLayout<Int32>.size))
        var address = sockaddr_in(); address.sin_len = UInt8(MemoryLayout<sockaddr_in>.size)
        address.sin_family = sa_family_t(AF_INET); address.sin_port = Self.port.bigEndian
        address.sin_addr = in_addr(s_addr: inet_addr("127.0.0.1"))
        let result = withUnsafePointer(to: &address) { pointer in
            pointer.withMemoryRebound(to: sockaddr.self, capacity: 1) { Darwin.bind(fd, $0, socklen_t(MemoryLayout<sockaddr_in>.size)) }
        }
        guard result == 0 else { let code = errno; Darwin.close(fd); throw POSIXError(.init(rawValue: code) ?? .EIO) }
        return fd
    }
    private func udpLoop() {
        while running {
            var bytes = [UInt8](repeating: 0, count: 4096); var peer = sockaddr_storage(); var length = socklen_t(MemoryLayout.size(ofValue: peer))
            let count = withUnsafeMutablePointer(to: &peer) { pointer in
                pointer.withMemoryRebound(to: sockaddr.self, capacity: 1) { recvfrom(udp, &bytes, bytes.count, 0, $0, &length) }
            }
            guard count > 0 else { if errno == EINTR { continue }; return }
            guard let response = try? MeshDNSCodec.response(to: Data(bytes.prefix(count)), aliases: aliases()) else { continue }
            response.withUnsafeBytes { body in
                withUnsafePointer(to: &peer) { pointer in
                    pointer.withMemoryRebound(to: sockaddr.self, capacity: 1) { _ = sendto(udp, body.baseAddress, response.count, 0, $0, length) }
                }
            }
        }
    }
    private func tcpLoop() {
        while running {
            let client = Darwin.accept(tcp, nil, nil)
            if client < 0 { if errno == EINTR { continue }; return }
            DispatchQueue.global(qos: .utility).async { [weak self] in self?.handleTCP(client) }
        }
    }
    private func handleTCP(_ fd: Int32) {
        defer { Darwin.close(fd) }
        guard let prefix = readExact(2, fd), prefix.count == 2 else { return }
        let length = Int(prefix[0]) << 8 | Int(prefix[1]); guard length > 0, length <= 4096,
              let query = readExact(length, fd), let response = try? MeshDNSCodec.response(to: query, aliases: aliases()) else { return }
        var framed = Data([UInt8(response.count >> 8), UInt8(response.count & 255)]); framed.append(response)
        framed.withUnsafeBytes { bytes in
            guard let base = bytes.baseAddress else { return }
            var offset = 0
            while offset < framed.count {
                let count = Darwin.write(fd, base.advanced(by: offset), framed.count - offset)
                if count <= 0 { return }
                offset += count
            }
        }
    }
    private func readExact(_ count: Int, _ fd: Int32) -> Data? {
        var output = Data(); var buffer = [UInt8](repeating: 0, count: count)
        while output.count < count {
            let n = Darwin.read(fd, &buffer, count - output.count); if n <= 0 { return nil }
            output.append(buffer, count: n)
        }
        return output
    }
}