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

import Foundation
import NIOCore
import NIOPosix

// 处理和sn-server服务器之间的通讯
final class SDLDNSClient: ChannelInboundHandler {
    typealias InboundIn = AddressedEnvelope<ByteBuffer>
    
    private let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
    
    private var channel: Channel?
    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)
    }
    
    func start() throws {
        let bootstrap = DatagramBootstrap(group: group)
            .channelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1)
            .channelInitializer { channel in
                channel.pipeline.addHandler(self)
            }
        
        let channel = try bootstrap.bind(host: "0.0.0.0", port: 0).wait()
        self.logger.log("[DNSClient] started", level: .debug)
        self.channel = channel
    }
    
    func waitClose() async throws {
        try await self.channel?.closeFuture.get()
    }
    
    // --MARK: ChannelInboundHandler delegate
    
    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
        let envelope = unwrapInboundIn(data)
        
        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))
        }
    }
    
    func channelInactive(context: ChannelHandlerContext) {
        self.packetContinuation.finish()
    }
        
    func forward(ipPacket: IPPacket) {
        guard let channel = self.channel else {
            return
        }
        
        let buffer = channel.allocator.buffer(bytes: ipPacket.data)
        let envelope = AddressedEnvelope<ByteBuffer>(remoteAddress: self.dnsServerAddress, data: buffer)
        channel.pipeline.eventLoop.execute {
            channel.writeAndFlush(envelope, promise: nil)
        }
    }
    
    deinit {
        try? self.group.syncShutdownGracefully()
        self.packetContinuation.finish()
    }
    
}

extension SDLDNSClient {
    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
        }
    }
    
}