import Foundation
import NetworkExtension

actor DNSService {
    enum Event {
        case packet(Data)
    }

    typealias EventHandler = @Sendable (Event) async -> Void

    private let serverIP: String
    private let publicDnsServers: [String]
    private let onEvent: EventHandler

    private var dnsClient: DNSCloudClient?
    private var dnsMonitorTask: Task<Void, Never>?

    private var dnsLocalClient: DNSLocalClient?
    private var dnsLocalMonitorTask: Task<Void, Never>?

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

    func start() {
        self.startCloud()
        self.startLocal()
    }

    func stop() async {
        let dnsClient = self.dnsClient
        self.dnsClient = nil
        let dnsMonitorTask = self.dnsMonitorTask
        self.dnsMonitorTask = nil

        let dnsLocalClient = self.dnsLocalClient
        self.dnsLocalClient = nil
        let dnsLocalMonitorTask = self.dnsLocalMonitorTask
        self.dnsLocalMonitorTask = nil

        dnsMonitorTask?.cancel()
        dnsClient?.stop()
        if let dnsMonitorTask {
            await dnsMonitorTask.value
        }

        dnsLocalMonitorTask?.cancel()
        await dnsLocalClient?.stop()
        if let dnsLocalMonitorTask {
            await dnsLocalMonitorTask.value
        }
    }

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

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

    private func startCloud() {
        guard self.dnsMonitorTask == nil else {
            return
        }

        self.dnsMonitorTask = startMonitorTask(name: "dnsServiceCloudMonitor") { [weak self] in
            guard let self else {
                throw CancellationError()
            }
            try await self.runCloud()
        }
    }

    private func startLocal() {
        guard self.dnsLocalMonitorTask == nil else {
            return
        }

        self.dnsLocalMonitorTask = startMonitorTask(name: "dnsServiceLocalMonitor") { [weak self] in
            guard let self else {
                throw CancellationError()
            }
            try await self.runLocal()
        }
    }

    private func runCloud() async throws {
        let dnsClient = DNSCloudClient(serverIP: self.serverIP, port: 15353)
        self.dnsClient = dnsClient
        dnsClient.start()

        defer {
            dnsClient.stop()
            if self.dnsClient === dnsClient {
                self.dnsClient = nil
            }
        }

        let onEvent = self.onEvent
        try await withTaskCancellationHandler {
            for try await packet in dnsClient.packetFlow {
                try Task.checkCancellation()
                await onEvent(.packet(packet))
            }
        } onCancel: {
            dnsClient.stop()
        }
    }

    private func runLocal() async throws {
        let dnsServer = self.publicDnsServers.randomElement() ?? "223.5.5.5"
        let dnsLocalClient = DNSLocalClient(host: dnsServer)
        await dnsLocalClient.start()
        self.dnsLocalClient = dnsLocalClient
        SDLLogger.log("[DNSService] dnsLocalClient started")

        defer {
            if self.dnsLocalClient === dnsLocalClient {
                self.dnsLocalClient = nil
            }
        }

        let onEvent = self.onEvent
        do {
            try await withTaskCancellationHandler {
                for try await packet in dnsLocalClient.packetFlow {
                    try Task.checkCancellation()
                    await onEvent(.packet(packet))
                }
            } onCancel: {
                Task {
                    await dnsLocalClient.stop()
                }
            }
            await dnsLocalClient.stop()
        } catch {
            await dnsLocalClient.stop()
            throw error
        }
    }
}