//
//  DNSClient.swift
//  Tun
//
//  Created by 安礼成 on 2025/12/10.
//

import Foundation
import NIOCore
import NIOPosix

// 处理和sn-server服务器之间的通讯
@available(macOS 14, *)
actor SDLDNSClientActor {
    private let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
    private let asyncChannel: NIOAsyncChannel<AddressedEnvelope<ByteBuffer>, AddressedEnvelope<ByteBuffer>>
    private let (writeStream, writeContinuation) = AsyncStream.makeStream(of: Data.self, bufferingPolicy: .unbounded)
    
    private let logger: SDLLogger
    private let dnsServerAddress: SocketAddress
    
    public let packetFlow: AsyncStream<Data>
    private let packetContinuation: AsyncStream<Data>.Continuation

    // 启动函数
    init(dnsServerAddress: SocketAddress, logger: SDLLogger) async throws {
        self.dnsServerAddress = dnsServerAddress
        self.logger = logger
        
        (self.packetFlow, self.packetContinuation) = AsyncStream.makeStream(of: Data.self, bufferingPolicy: .unbounded)
        
        let bootstrap = DatagramBootstrap(group: group)
            .channelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1)
        
        self.asyncChannel = try await bootstrap.bind(host: "0.0.0.0", port: 0)
            .flatMapThrowing { channel in
                return try NIOAsyncChannel(wrappingChannelSynchronously: channel, configuration: .init(
                    inboundType: AddressedEnvelope<ByteBuffer>.self,
                    outboundType: AddressedEnvelope<ByteBuffer>.self
                ))
            }
            .get()
    }
    
    func start() async throws {
        try await withTaskCancellationHandler {
            try await self.asyncChannel.executeThenClose {inbound, outbound in
                try await withThrowingTaskGroup(of: Void.self) { group in
                    group.addTask {
                        defer {
                            self.logger.log("[DNSClient] inbound closed", level: .warning)
                        }
                        
                        for try await envelope in inbound {
                            try Task.checkCancellation()
                            var buffer = envelope.data
                            let remoteAddress = envelope.remoteAddress
                            self.logger.log("[DNSClient] read data: \(buffer), from: \(remoteAddress)", level: .debug)
                            
                            let len = buffer.readableBytes
                            if let bytes = buffer.readBytes(length: len) {
                                self.packetContinuation.yield(Data(bytes))
                            }
                        }
                    }
                    
                    group.addTask {
                        defer {
                            self.logger.log("[DNSClient] outbound closed", level: .warning)
                        }
                        
                        for await message in self.writeStream {
                            try Task.checkCancellation()
                            
                            let buffer = self.asyncChannel.channel.allocator.buffer(bytes: message)
                            let envelope = AddressedEnvelope<ByteBuffer>(remoteAddress: self.dnsServerAddress, data: buffer)
                            try await outbound.write(envelope)
                        }
                    }
                    
                    if let _ = try await group.next() {
                        group.cancelAll()
                    }
                }
            }
        } onCancel: {
            self.writeContinuation.finish()
            self.packetContinuation.finish()
            self.logger.log("[DNSClient] withTaskCancellationHandler cancel")
        }
    }
    
    func forward(ipPacket: IPPacket) {
        self.writeContinuation.yield(ipPacket.data)
    }
    
    deinit {
        try? self.group.syncShutdownGracefully()
        self.writeContinuation.finish()
    }
    
}

extension SDLDNSClientActor {
    
    struct Helper {
        static let dnsServer: String = "100.100.100.100"
        // dns请求包的目标地址
        static let dnsDestIpAddr: UInt32 = 1684300900
        
        // 判断是否是dns请求的数据包
        static func isDnsRequestPacket(ipPacket: IPPacket) -> Bool {
            return ipPacket.header.destination == dnsDestIpAddr
        }
    }
    
}