//
//  File.swift
//  sdlan
//
//  Created by 安礼成 on 2025/7/14.
//

import Foundation
import NIOCore

// 网络类型探测器
@available(macOS 14, *)
struct SDLNatProber {
    
    // 定义nat类型
    enum NatType: UInt8, Encodable {
        case blocked = 0
        case noNat = 1
        case fullCone = 2
        case portRestricted = 3
        case coneRestricted = 4
        case symmetric = 5
    }
    
    // 获取当前所处的网络的nat类型
    static func getNatType(udpHole: SDLUDPHole?, config: SDLConfiguration, logger: SDLLogger) async -> NatType {
        guard let udpHole else {
            return .blocked
        }
        
        let addressArray = config.stunProbeSocketAddressArray
        // step1: ip1:port1 <---- ip1:port1
        guard let natAddress1 = await getNatAddress(udpHole, remoteAddress: addressArray[0][0], attr: .none) else {
            return .blocked
        }
        
        // 网络没有在nat下
        if await natAddress1 == udpHole.localAddress {
            return .noNat
        }
        
        // step2: ip2:port2 <---- ip2:port2
        guard let natAddress2 = await getNatAddress(udpHole, remoteAddress: addressArray[1][1], attr: .none) else {
            return .blocked
        }
        
        // 如果natAddress2 的IP地址与上次回来的IP是不一样的，它就是对称型NAT; 这次的包也一定能发成功并收到
        // 如果ip地址变了，这说明{dstIp, dstPort, srcIp, srcPort}, 其中有一个变了；则用新的ip地址
        logger.log("[SDLNatProber] nat_address1: \(natAddress1), nat_address2: \(natAddress2)", level: .debug)
        if let ipAddress1 = natAddress1.ipAddress, let ipAddress2 = natAddress2.ipAddress, ipAddress1 != ipAddress2 {
            return .symmetric
        }
        
        // step3: ip1:port1 <---- ip2:port2 (ip地址和port都变的情况)
        // 如果能收到的，说明是完全锥形 说明是IP地址限制锥型NAT，如果不能收到说明是端口限制锥型。
        if let natAddress3 = await getNatAddress(udpHole, remoteAddress: addressArray[0][0], attr: .peer) {
            logger.log("[SDLNatProber] nat_address1: \(natAddress1), nat_address2: \(natAddress2), nat_address3: \(natAddress3)", level: .debug)
            return .fullCone
        }
        
        // step3: ip1:port1 <---- ip1:port2 (port改变情况)
        // 如果能收到的说明是IP地址限制锥型NAT，如果不能收到说明是端口限制锥型。
        if let natAddress4 = await getNatAddress(udpHole, remoteAddress: addressArray[0][0], attr: .port) {
            logger.log("[SDLNatProber] nat_address1: \(natAddress1), nat_address2: \(natAddress2), nat_address4: \(natAddress4)", level: .debug)
            return .coneRestricted
        } else {
            return .portRestricted
        }
    }
    
    private static func getNatAddress(_ udpHole: SDLUDPHole, remoteAddress: SocketAddress, attr: SDLProbeAttr) async -> SocketAddress? {
        let stunProbeReply = try? await udpHole.stunProbe(remoteAddress: remoteAddress, attr: attr, timeout: 5)
        
        return stunProbeReply?.socketAddress()
    }
    
}