//
//  UDPSocksToken.swift
//  MpAccSocks5SDK
//

import Foundation

// 参考: https://www.jianshu.com/p/cf88c619ee5c
struct DatagramFrame {
    var frag: UInt8
    var dstAddr: UInt32
    var dstPort: UInt16
    var data: Data
    func getFrameData() -> Data {
        var frameData = Data([
            0x00, 0x00, // 保留位为0
            frag,   // 片位
            0x01,   // 地址类型，必须为ipv4
            // ip 地址
            UInt8(dstAddr & 0xFF),
            UInt8(dstAddr >> 8 & 0xFF),
            UInt8(dstAddr >> 16 & 0xFF),
            UInt8(dstAddr >> 24 & 0xFF),
            // port 端口
            UInt8(dstPort & 0xFF),
            UInt8(dstPort >> 8 & 0xFF),
        ])
        frameData.append(data)
        return frameData
    }
    
    static func frame(from frameData: Data) -> DatagramFrame? {
        guard frameData.count >= 10 else { return nil }
        let bytes = [UInt8](frameData)
        guard bytes[0] == 0, bytes[1] == 0,    // 保留字段为0
              bytes[3] == 0x01  // 地址类型必须为ipv4
        else {
            return nil
        }
        return DatagramFrame(
            frag: bytes[2],
            dstAddr: (UInt32(bytes[4]) | UInt32(bytes[5]) << 8 | UInt32(bytes[6]) << 16 | UInt32(bytes[7]) << 24),
            dstPort: (UInt16(bytes[8]) | UInt16(bytes[9]) << 8),
            data: Data(bytes[10..<bytes.count])
        )
    }
}

public class UDPSocksToken: NSObject {
    private var socketFd: Int32
    @objc public private(set) var disposed = false
    @objc public let proxyIp: String
    @objc public let proxyPort: UInt16
    
    fileprivate let dstAddr: UInt32
    fileprivate let dstPort: UInt16
    
    
    init(socketFd: Int32, proxyIp: String, proxyPort: UInt16, dstAddr: UInt32, dstPort: UInt16) {
        assert(socketFd != -1)
        self.socketFd = socketFd
        self.proxyIp = proxyIp
        self.proxyPort = proxyPort
        self.dstAddr = dstAddr
        self.dstPort = dstPort
    }
    
    @objc
    public func createUdpFrame(frag:UInt8 = 0x00, orgData: Data) -> Data {
        guard disposed == false else { return orgData }
        let frame = DatagramFrame(frag: frag, dstAddr: dstAddr, dstPort: dstPort, data: orgData)
        return frame.getFrameData()
    }
    
    @objc
    public func getOrgDataFromDatagram(_ data: Data, from ip: String, port: UInt16) -> Data {
        guard disposed == false else { return data }
        guard ip == proxyIp, port == proxyPort, data.count > 10 else { return data }
        guard let frame = DatagramFrame.frame(from: data) else { return data }
        guard frame.dstAddr == dstAddr,
              frame.dstPort == dstPort
        else {
            return data
        }
        return frame.data
    }
    
    @objc
    public func dispose() {
        if !disposed && socketFd != -1 {
            close(socketFd)
        }
        disposed = true
        socketFd = -1
    }
    
    deinit {
        if socketFd != -1 {
            close(socketFd)
        }
    }
}
