//
//  SDLDNSClient 2.swift
//  punchnet
//
//  Created by 安礼成 on 2026/4/9.
//
import Foundation
import Network

final class DNSCloudClient {
    private var connection: NWConnection?
    private let logger: SDLLogger
    private let dnsServerAddress: NWEndpoint
    
    // 用于对外输出收到的 DNS 响应包
    public let packetFlow: AsyncStream<Data>
    private let packetContinuation: AsyncStream<Data>.Continuation
    
    // 用来处理关闭事件
    private let (closeStream, closeContinuation) = AsyncStream.makeStream(of: Void.self)
    
    /// - Parameter host: 你的 sn-server 地址 (如 "8.8.8.8")
    /// - Parameter port: 端口 (如 53)
    init(host: String, port: UInt16, logger: SDLLogger) {
        self.logger = logger
        self.dnsServerAddress = .hostPort(host: NWEndpoint.Host(host), port: NWEndpoint.Port(integerLiteral: port))
        
        let (stream, continuation) = AsyncStream.makeStream(of: Data.self, bufferingPolicy: .unbounded)
        self.packetFlow = stream
        self.packetContinuation = continuation
    }
    
    func start() {
        // 1. 配置参数：这是解决环路的关键
        let parameters = NWParameters.udp
        
        // 禁止此连接走 TUN 网卡（在 NE 中 TUN 通常被归类为 .other）
        parameters.prohibitedInterfaceTypes = [.other]
        // 2. 增强健壮性：启用多路径切换（替代 pathSelectionOptions 的意图）
        parameters.multipathServiceType = .handover
        
        // 2. 创建连接
        let connection = NWConnection(to: self.dnsServerAddress, using: parameters)
        self.connection = connection
        
        connection.stateUpdateHandler = { [weak self] state in
            switch state {
            case .ready:
                self?.logger.log("[DNSClient] Connection ready", level: .debug)
                self?.receiveLoop() // 开始循环接收数据
            case .failed(let error):
                self?.logger.log("[DNSClient] Connection failed: \(error)", level: .error)
                self?.stop()
            case .cancelled:
                self?.packetContinuation.finish()
                self?.closeContinuation.finish()
            default:
                break
            }
        }
        
        // 启动连接队列
        connection.start(queue: .global())
    }
    
    public func waitClose() async {
        for await _ in closeStream { }
    }
    
    /// 接收数据的递归循环
    private func receiveLoop() {
        connection?.receiveMessage { [weak self] content, _, isComplete, error in
            if let data = content, !data.isEmpty {
                // 将收到的 DNS 响应写回 AsyncStream
                self?.packetContinuation.yield(data)
            }
            
            if error == nil && self?.connection?.state == .ready {
                self?.receiveLoop() // 继续监听下一个包
            }
        }
    }
    
    /// 发送 DNS 查询包（由 TUN 拦截到的原始 IP 包数据）
    func forward(ipPacketData: Data) {
        guard let connection = self.connection, connection.state == .ready else {
            return
        }
        
        connection.send(content: ipPacketData, completion: .contentProcessed { [weak self] error in
            if let error = error {
                self?.logger.log("[DNSClient] Send error: \(error)", level: .error)
            }
        })
    }
    
    func stop() {
        connection?.cancel()
        connection = nil
    }
    
    deinit {
        stop()
    }
    
}