import Foundation

actor DNSLocalService {
    private let publicDnsServers: [String]
    private var onEvent: DNSEventHandler = { _ in }

    private var currentClient: DNSLocalClient?
    private var generation: UInt64 = 0

    init(publicDnsServers: [String]) {
        self.publicDnsServers = publicDnsServers
    }

    func updateEventHandler(_ onEvent: @escaping DNSEventHandler) {
        self.onEvent = onEvent
    }

    func run() async throws {
        let generation = self.nextGeneration()
        let dnsServer = self.publicDnsServers.randomElement() ?? "223.5.5.5"
        let client = DNSLocalClient(host: dnsServer)
        self.currentClient = client

        SDLLogger.log("[DNSLocalService] dnsLocalClient started")

        do {
            try await self.run(client: client)
            self.clearCurrent(client, generation: generation)
            await client.stop()
        } catch is CancellationError {
            self.clearCurrent(client, generation: generation)
            await client.stop()
            throw CancellationError()
        } catch {
            self.clearCurrent(client, generation: generation)
            await client.stop()
            throw error
        }
    }

    func stop() async {
        self.generation &+= 1

        let client = self.currentClient
        self.currentClient = nil

        await client?.stop()
    }

    func query(tracker: DNSLocalClient.DNSTracker, dnsPayload: Data) async {
        await self.currentClient?.query(tracker: tracker, dnsPayload: dnsPayload)
    }

    private func run(client: DNSLocalClient) async throws {
        let onEvent = self.onEvent

        try await withThrowingTaskGroup(of: Void.self) { group in
            defer {
                group.cancelAll()
            }

            group.addTask {
                try await client.run()
            }

            group.addTask {
                for try await packet in client.packetFlow {
                    try Task.checkCancellation()
                    await onEvent(.packet(packet))
                }
            }

            _ = try await group.next()
        }
    }

    private func nextGeneration() -> UInt64 {
        self.generation &+= 1
        return self.generation
    }

    private func clearCurrent(_ client: DNSLocalClient, generation: UInt64) {
        guard self.generation == generation else {
            return
        }

        if self.currentClient === client {
            self.currentClient = nil
        }
    }
}