import Foundation

actor DNSCloudService {
    private let serverIP: String
    private let onEvent: DNSService.EventHandler

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

    init(serverIP: String, onEvent: @escaping DNSService.EventHandler) {
        self.serverIP = serverIP
        self.onEvent = onEvent
    }

    func run() async throws {
        let generation = self.nextGeneration()
        let client = DNSCloudClient(serverIP: self.serverIP, port: 15353)
        self.currentClient = client

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

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

        let client = self.currentClient
        self.currentClient = nil

        client?.stop()
    }

    func forward(ipPacketData: Data) {
        self.currentClient?.forward(ipPacketData: ipPacketData)
    }

    private func run(client: DNSCloudClient) 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: DNSCloudClient, generation: UInt64) {
        guard self.generation == generation else {
            return
        }

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