import Foundation
import AsyncDNSResolver

struct TimeoutError: Error {}

func getSRVRecords(withTimeout: TimeInterval) async throws -> [SRVRecord] {
    let resolver = try AsyncDNSResolver()
    let records = try await withThrowingTaskGroup(of: [SRVRecord].self) { group in
        group.addTask {
            return try await resolver.querySRV(name: "")
        }
        group.addTask {
            try await Task.sleep(nanoseconds: UInt64(withTimeout) * NSEC_PER_SEC)
            try Task.checkCancellation()
            throw TimeoutError()
        }
        do {
            let records = try await group.next()!
            group.cancelAll()
            return records
        } catch {
            group.cancelAll()
            throw error
        }
    }
    return records
}

do {
    let records = try await getSRVRecords(withTimeout: 3)
    print(records)
} catch {
    print("Error: \(error)")
}

// Source: https://stackoverflow.com/questions/22502729/resolving-srv-records-with-ios-sdk

- (void)updateDnsRecords
{
    if (self.dnsUpdatePending == YES)
    {
        return;
    }
    else
    {
        self.dnsUpdatePending = YES;
    }

    NSLog(@"DNS update");
    DNSServiceRef       sdRef;
    DNSServiceErrorType err;

    const char* host = [self.dnsHost UTF8String];
    if (host != NULL)
    {
        NSTimeInterval remainingTime = self.dnsUpdateTimeout;
        NSDate*        startTime = [NSDate date];

        err = DNSServiceQueryRecord(&sdRef, 0, 0,
                                    host,
                                    kDNSServiceType_SRV,
                                    kDNSServiceClass_IN,
                                    processDnsReply,
                                    &remainingTime);

        // This is necessary so we don't hang forever if there are no results
        int            dns_sd_fd = DNSServiceRefSockFD(sdRef);
        int            nfds      = dns_sd_fd + 1;
        fd_set         readfds;
        int            result;

        while (remainingTime > 0)
        {
            FD_ZERO(&readfds);
            FD_SET(dns_sd_fd, &readfds);

            struct timeval tv;
            tv.tv_sec  = (time_t)remainingTime;
            tv.tv_usec = (remainingTime - tv.tv_sec) * 1000000;

            result = select(nfds, &readfds, (fd_set*)NULL, (fd_set*)NULL, &tv);
            if (result == 1)
            {
                if (FD_ISSET(dns_sd_fd, &readfds))
                {
                    err = DNSServiceProcessResult(sdRef);
                    if (err != kDNSServiceErr_NoError)
                    {
                        NSLog(@"There was an error reading the DNS SRV records.");
                        break;
                    }
                }
            }
            else if (result == 0)
            {
                NBLog(@"DNS SRV select() timed out");
                break;
            }
            else
            {
                if (errno == EINTR)
                {
                    NBLog(@"DNS SRV select() interrupted, retry.");
                }
                else
                {
                    NBLog(@"DNS SRV select() returned %d errno %d %s.", result, errno, strerror(errno));
                    break;
                }
            }

            NSTimeInterval elapsed = [[NSDate date] timeIntervalSinceDate:startTime];
            remainingTime -= elapsed;
        }

        DNSServiceRefDeallocate(sdRef);
    }
}