import Foundation
import NetworkExtension

actor SDLPacketReaderService {
    enum Event {
        case packet(IPPacket)
    }

    typealias EventHandler = @Sendable (Event) async -> Void
    private typealias PacketReadResult = (packets: [Data], protocols: [NSNumber])?

    private final class CancellationToken: @unchecked Sendable {
        private let lock = NSLock()
        private var cancelled: Bool = false

        var isCancelled: Bool {
            lock.lock()
            defer {
                lock.unlock()
            }
            return cancelled
        }

        func cancel() {
            lock.lock()
            cancelled = true
            lock.unlock()
        }
    }

    private let provider: NEPacketTunnelProvider
    private let onEvent: EventHandler
    private var readTask: Task<Void, Never>?
    private var readToken: CancellationToken?

    init(provider: NEPacketTunnelProvider, onEvent: @escaping EventHandler) {
        self.provider = provider
        self.onEvent = onEvent
    }

    func start() {
        guard self.readTask == nil else {
            return
        }

        let provider = self.provider
        let onEvent = self.onEvent
        let readToken = CancellationToken()
        self.readToken = readToken

        self.readTask = Task(priority: .high) {
            while !Task.isCancelled && !readToken.isCancelled {
                guard let batch = await Self.readPackets(from: provider) else {
                    break
                }

                if Task.isCancelled || readToken.isCancelled {
                    break
                }

                for (data, number) in zip(batch.packets, batch.protocols) where number.int32Value == 2 {
                    if Task.isCancelled || readToken.isCancelled {
                        break
                    }

                    if let packet = IPPacket(data) {
                        await onEvent(.packet(packet))
                    }
                }
            }
            SDLLogger.log("[SDLPacketReaderService] readTask finished")
        }
    }

    func stop() {
        let readTask = self.readTask
        self.readTask = nil
        let readToken = self.readToken
        self.readToken = nil
        readToken?.cancel()
        readTask?.cancel()
    }

    private static func readPackets(from provider: NEPacketTunnelProvider) async -> PacketReadResult {
        let readContinuation = OnceContinuation<PacketReadResult, Never>()

        return await withTaskCancellationHandler {
            await withCheckedContinuation { continuation in
                readContinuation.set(continuation)
                provider.packetFlow.readPackets { packets, protocols in
                    readContinuation.resume(returning: (packets: packets, protocols: protocols))
                }
            }
        } onCancel: {
            readContinuation.resume(returning: nil)
        }
    }
}