//
//  SDLContext.swift
//  Tun
//
//  Created by 安礼成 on 2024/2/29.
//

import Foundation
import NetworkExtension
import NIOCore
import Combine

// 上下文环境变量，全局共享
/*
1. 处理rsa的加解密逻辑
 */

public class SDLContext: @unchecked Sendable {
    
    // 路由信息
    struct Route {
        let dstAddress: String
        let subnetMask: String
        
        var debugInfo: String {
            return "\(dstAddress):\(subnetMask)"
        }
    }
    
    let config: SDLConfiguration

    // tun网络地址信息
    var devAddr: SDLDevAddr
    
    // nat映射的相关信息, 暂时没有用处
    //var natAddress: SDLNatAddress?
    // nat的网络类型
    var natType: SDLNatProber.NatType = .blocked
    
    // AES加密，授权通过后，对象才会被创建
    var aesCipher: AESCipher
    
    // aes
    var aesKey: Data = Data()
    
    // rsa的相关配置, public_key是本地生成的
    let rsaCipher: RSACipher
    
    // 依赖的变量
    var udpHoleActor: SDLUDPHoleActor?
    var superClientActor: SDLSuperClientActor?
    
    // 数据包读取任务
    private var readTask: Task<(), Never>?
    
    let provider: NEPacketTunnelProvider
    
    private var sessionManager: SessionManager
    private var holerManager: HolerManager
    private var arpServer: ArpServer
    
    // 记录最后发送的stunRequest的cookie
    private var lastCookie: UInt32? = 0
    
    // 定时器
    private var stunCancel: AnyCancellable?
    
    // 网络状态变化的健康
    private var monitor = SDLNetworkMonitor()
    private var monitorCancel: AnyCancellable?
    
    // 内部socket通讯
    private var noticeClient: SDLNoticeClient
    
    // 流量统计
    private var flowTracer = SDLFlowTracerActor()
    private var flowTracerCancel: AnyCancellable?
    
    public init(provider: NEPacketTunnelProvider, config: SDLConfiguration, rsaCipher: RSACipher, aesCipher: AESCipher) {
        self.config = config
        self.rsaCipher = rsaCipher
        self.aesCipher = aesCipher
        
        // 生成mac地址
        var devAddr = SDLDevAddr()
        devAddr.mac = Self.getMacAddress()
        self.devAddr = devAddr
        
        self.provider = provider
        self.sessionManager = SessionManager()
        self.holerManager = HolerManager()
        self.arpServer = ArpServer(known_macs: [:])
        self.noticeClient = SDLNoticeClient()
    }
    
    public func start() async throws {
        self.udpHoleActor = try await SDLUDPHoleActor()
        self.superClientActor = try await SDLSuperClientActor(host: self.config.superHost, port: self.config.superPort)
        
        try await withThrowingTaskGroup(of: Void.self) { group in
            group.addTask {
                try await self.udpHoleActor?.start()
            }
            
            group.addTask {
                try await self.superClientActor?.start()
            }
            
            group.addTask {
                if let eventFlow = self.superClientActor?.eventFlow {
                    for try await event in eventFlow {
                        try await self.handleSuperEvent(event: event)
                    }
                }
            }
            
            group.addTask {
                if let eventFlow = self.udpHoleActor?.eventFlow {
                    for try await event in eventFlow {
                        try await self.handleUDPEvent(event: event)
                    }
                }
                
            }
            
            try await group.waitForAll()
        }
        
//        self.noticeClient.start()
//        // 启动网络监控
//        self.monitorCancel = self.monitor.eventFlow.sink { event in
//            switch event {
//            case .changed:
//                // 需要重新探测网络的nat类型
//                Task {
//                    self.natType = await SDLNatProber.getNatType(udpHole: self.udpHole, config: self.config)
//                    NSLog("didNetworkPathChanged, nat type is: \(self.natType)")
//                }
//            case .unreachable:
//                NSLog("didNetworkPathUnreachable")
//            }
//        }
//        self.monitor.start()
    }
    
    public func stop() async {
        self.superClientActor = nil
        self.udpHoleActor = nil
        
        self.readTask?.cancel()
    }
    
    private func handleSuperEvent(event: SDLSuperClientActor.SuperEvent) async throws {
        switch event {
        case .ready:
            NSLog("[SDLContext] get registerSuper, mac address: \(Self.formatMacAddress(mac: self.devAddr.mac))")
            guard let message = try await self.superClientActor?.registerSuper(context: self) else {
                return
            }
            
            switch message.packet {
            case .registerSuperAck(let registerSuperAck):
                // 需要对数据通过rsa的私钥解码
                let aesKey = try! self.rsaCipher.decode(data: Data(registerSuperAck.aesKey))
                let upgradeType = SDLUpgradeType(rawValue: registerSuperAck.upgradeType)
                
                NSLog("[SDLContext] get registerSuperAck, aes_key len: \(aesKey.count)")
                self.devAddr = registerSuperAck.devAddr
                
                if upgradeType == .force {
                    let forceUpgrade = NoticeMessage.UpgradeMessage(prompt: registerSuperAck.upgradePrompt, address: registerSuperAck.upgradeAddress)
                    self.noticeClient.send(data: forceUpgrade.binaryData)
                    exit(-1)
                }
                
                // 服务器分配的tun网卡信息
                await self.didNetworkConfigChanged(devAddr: self.devAddr)
                self.aesKey = aesKey
                
                if upgradeType == .normal {
                    let normalUpgrade = NoticeMessage.UpgradeMessage(prompt: registerSuperAck.upgradePrompt, address: registerSuperAck.upgradeAddress)
                    self.noticeClient.send(data: normalUpgrade.binaryData)
                }
                
            case .registerSuperNak(let nakPacket):
                let errorMessage = nakPacket.errorMessage
                guard let errorCode = SDLNAKErrorCode(rawValue: UInt8(nakPacket.errorCode)) else {
                    return
                }
                
                switch errorCode {
                case .invalidToken, .nodeDisabled:
                    let alertNotice = NoticeMessage.AlertMessage(alert: errorMessage)
                    self.noticeClient.send(data: alertNotice.binaryData)
                    exit(-1)
                case .noIpAddress, .networkFault, .internalFault:
                    let alertNotice = NoticeMessage.AlertMessage(alert: errorMessage)
                    self.noticeClient.send(data: alertNotice.binaryData)
                }
                NSLog("[SDLContext] Get a SuperNak message exit")
            default:
                ()
            }
            
        case .closed:
            NSLog("[SDLContext] super client closed")
            await self.arpServer.clear()
            DispatchQueue.main.asyncAfter(deadline: .now() + 5) {
//                Task {@MainActor in
//                    try await self.startSuperClient()
//                }
            }
        case .event(let evt):
            switch evt {
            case .natChanged(let natChangedEvent):
                let dstMac = natChangedEvent.mac
                NSLog("[SDLContext] natChangedEvent, dstMac: \(dstMac)")
                await sessionManager.removeSession(dstMac: dstMac)
            case .sendRegister(let sendRegisterEvent):
                NSLog("[SDLContext] sendRegisterEvent, ip: \(sendRegisterEvent)")
                let address = SDLUtil.int32ToIp(sendRegisterEvent.natIp)
                if let remoteAddress = try? SocketAddress.makeAddressResolvingHost(address, port: Int(sendRegisterEvent.natPort)) {
                    // 发送register包
                    await self.udpHoleActor?.sendRegister(context: self, remoteAddress: remoteAddress, dst_mac: sendRegisterEvent.dstMac)
                }
                
            case .networkShutdown(let shutdownEvent):
                let alertNotice = NoticeMessage.AlertMessage(alert: shutdownEvent.message)
                self.noticeClient.send(data: alertNotice.binaryData)
                exit(-1)
            }
        case .command(let packetId, let command):
            switch command {
            case .changeNetwork(let changeNetworkCommand):
                // 需要对数据通过rsa的私钥解码
                let aesKey = try! self.rsaCipher.decode(data: Data(changeNetworkCommand.aesKey))
                NSLog("[SDLContext] change network command get aes_key len: \(aesKey.count)")
                self.devAddr = changeNetworkCommand.devAddr
                
                // 服务器分配的tun网卡信息
                await self.didNetworkConfigChanged(devAddr: self.devAddr)
                self.aesKey = aesKey
                
                var commandAck = SDLCommandAck()
                commandAck.status = true
                
                await self.superClientActor?.commandAck(packetId: packetId, ack: commandAck)
            }
        }
        
    }

    private func startUDPHole() async throws {
//        self.udpHole = SDLUDPHole()
//        
//        self.udpCancel?.cancel()
//        self.udpCancel = self.udpHole?.eventFlow.sink { event in
//            Task.detached {
//                await self.handleUDPEvent(event: event)
//            }
//        }
//        
//        try await self.udpHole?.start()
    }
    
    private func handleUDPEvent(event: SDLUDPHoleActor.UDPEvent) async throws {
        switch event {
        case .ready:
            // 获取当前网络的类型
            self.natType = await SDLNatProber.getNatType(udpHole: self.udpHole, config: self.config)
            SDLLogger.log("[SDLContext] nat type is: \(self.natType)", level: .debug)
            
            let timer = Timer.publish(every: 5.0, on: .main, in: .common).autoconnect()
//            self.stunCancel = Just(Date()).merge(with: timer).sink { _ in
//                self.lastCookie = await self.udpHoleActor?.stunRequest(context: self)
//            }
            
        case .closed:
            DispatchQueue.main.asyncAfter(deadline: .now() + 5) {
                Task {
                    try await self.startUDPHole()
                }
            }
            
        case .message(let remoteAddress, let message):
            switch message {
            case .register(let register):
                NSLog("register packet: \(register), dev_addr: \(self.devAddr)")
                // 判断目标地址是否是tun的网卡地址, 并且是在同一个网络下
                if register.dstMac == self.devAddr.mac && register.networkID == self.devAddr.networkID {
                    // 回复ack包
                    self.udpHoleActor?.sendRegisterAck(context: self, remoteAddress: remoteAddress, dst_mac: register.srcMac)
                    // 这里需要建立到来源的会话, 在复杂网络下，通过super-node查询到的nat地址不一定靠谱，需要通过udp包的来源地址作为nat地址
                    let session = Session(dstMac: register.srcMac, natAddress: remoteAddress)
                    await self.sessionManager.addSession(session: session)
                } else {
                    SDLLogger.log("SDLContext didReadRegister get a invalid packet, because dst_ip not matched: \(register.dstMac)", level: .warning)
                }
            case .registerAck(let registerAck):
                // 判断目标地址是否是tun的网卡地址, 并且是在同一个网络下
                if registerAck.dstMac == self.devAddr.mac && registerAck.networkID == self.devAddr.networkID {
                    let session = Session(dstMac: registerAck.srcMac, natAddress: remoteAddress)
                    await self.sessionManager.addSession(session: session)
                } else {
                    SDLLogger.log("SDLContext didReadRegisterAck get a invalid packet, because dst_mac not matched: \(registerAck.dstMac)", level: .warning)
                }
            case .stunReply(let stunReply):
                let cookie = stunReply.cookie
                if cookie == self.lastCookie {
                    // 记录下当前在nat上的映射信息，暂时没有用；后续会用来判断网络类型
                    //self.natAddress = stunReply.natAddress
                    SDLLogger.log("[SDLContext] get a stunReply: \(try! stunReply.jsonString())")
                }
            default:
                ()
            }
            
        case .data(let data):
            let mac = LayerPacket.MacAddress(data: data.dstMac)
            guard (data.dstMac == self.devAddr.mac || mac.isBroadcast() || mac.isMulticast()) else {
                NSLog("[SDLContext] didReadData 1")
                return
            }
            
            guard let decyptedData = try? self.aesCipher.decypt(aesKey: self.aesKey, data: Data(data.data)) else {
                NSLog("[SDLContext] didReadData 2")
                return
            }
            
            do {
                let layerPacket = try LayerPacket(layerData: decyptedData)
                
                await self.flowTracer.inc(num: decyptedData.count, type: .inbound)
                // 处理arp请求
                switch layerPacket.type {
                case .arp:
                    // 判断如果收到的是arp请求
                    if let arpPacket = ARPPacket(data: layerPacket.data)  {
                        if arpPacket.targetIP == self.devAddr.netAddr {
                            switch arpPacket.opcode {
                            case .request:
                                NSLog("[SDLContext] get arp request packet")
                                let response = ARPPacket.arpResponse(for: arpPacket, mac: self.devAddr.mac, ip: self.devAddr.netAddr)
                                await self.routeLayerPacket(dstMac: arpPacket.senderMAC, type: .arp, data: response.marshal())
                            case .response:
                                NSLog("[SDLContext] get arp response packet")
                                await self.arpServer.append(ip: arpPacket.senderIP, mac: arpPacket.senderMAC)
                            }
                        } else {
                            NSLog("[SDLContext] get invalid arp packet, target_ip: \(arpPacket)")
                        }
                    } else {
                        NSLog("[SDLContext] get invalid arp packet")
                    }
                case .ipv4:
                    NSLog("[SDLContext] get ipv4 packet")
                    guard let ipPacket = IPPacket(layerPacket.data), ipPacket.header.destination == self.devAddr.netAddr else {
                        return
                    }
                    
                    let packet = NEPacket(data: ipPacket.data, protocolFamily: 2)
                    self.provider.packetFlow.writePacketObjects([packet])
                default:
                    NSLog("[SDLContext] get invalid packet")
                }
            } catch let err {
                NSLog("[SDLContext] didReadData err: \(err)")
            }
        }
        
    }
    
    // 流量统计
    public func flowReportTask() {
        Task {
            // 每分钟汇报一次
            self.flowTracerCancel = Timer.publish(every: 60.0, on: .main, in: .common).autoconnect()
                .sink { _ in
                    Task {
                        let (forwardNum, p2pNum, inboundNum) = await self.flowTracer.reset()
                        await self.superClientActor?.flowReport(forwardNum: forwardNum, p2pNum: p2pNum, inboundNum: inboundNum)
                    }
                }
        }
    }
    
    // 网络改变时需要重新配置网络信息
    private func didNetworkConfigChanged(devAddr: SDLDevAddr, dnsServers: [String]? = nil) async {
        let netAddress = SDLNetAddress(ip: devAddr.netAddr, maskLen: UInt8(devAddr.netBitLen))
        let routes = [Route(dstAddress: netAddress.networkAddress, subnetMask: netAddress.maskAddress)]
        
        // Add code here to start the process of connecting the tunnel.
        let networkSettings = NEPacketTunnelNetworkSettings(tunnelRemoteAddress: "8.8.8.8")
        networkSettings.mtu = 1460
        
        // 设置网卡的DNS解析
        if let dnsServers {
            networkSettings.dnsSettings = NEDNSSettings(servers: dnsServers)
        } else {
            networkSettings.dnsSettings = NEDNSSettings(servers: ["8.8.8.8", "114.114.114.114"])
        }
        
        NSLog("[SDLContext] Tun started at network ip: \(netAddress.ipAddress), mask: \(netAddress.maskAddress)")

        let ipv4Settings = NEIPv4Settings(addresses: [netAddress.ipAddress], subnetMasks: [netAddress.maskAddress])
        // 设置路由表
        //NEIPv4Route.default()
        ipv4Settings.includedRoutes = routes.map { route in
            NEIPv4Route(destinationAddress: route.dstAddress, subnetMask: route.subnetMask)
        }
        networkSettings.ipv4Settings = ipv4Settings
        // 网卡配置设置必须成功
        do {
            try await self.provider.setTunnelNetworkSettings(networkSettings)
            
            await self.holerManager.cleanup()
            self.startReader()
            
            NSLog("[SDLContext] setTunnelNetworkSettings success, start read packet")
        } catch let err {
            NSLog("[SDLContext] setTunnelNetworkSettings get error: \(err)")
            exit(-1)
        }
    }
    
    // 开始读取数据, 用单独的线程处理packetFlow
    private func startReader() {
        //  停止之前的任务
        self.readTask?.cancel()
        
        // 开启新的任务
        self.readTask = Task(priority: .high) {
            repeat {
                if Task.isCancelled {
                    break
                }
                
                let (packets, numbers) = await self.provider.packetFlow.readPackets()
                for (data, number) in zip(packets, numbers) where number == 2 {
                    if let packet = IPPacket(data) {
                        Task.detached {
                            let dstIp = packet.header.destination
                            // 本地通讯, 目标地址是本地服务器的ip地址
                            if dstIp == self.devAddr.netAddr {
                                let nePacket = NEPacket(data: packet.data, protocolFamily: 2)
                                self.provider.packetFlow.writePacketObjects([nePacket])
                                return
                            }
                            
                            // 查找arp缓存中是否有目标mac地址
                            if let dstMac = await self.arpServer.query(ip: dstIp) {
                                await self.routeLayerPacket(dstMac: dstMac, type: .ipv4, data: packet.data)
                            }
                            else {
                                // 构造arp请求
                                let broadcastMac = Data([0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF])
                                let arpReqeust: ARPPacket = ARPPacket.arpRequest(senderIP: self.devAddr.netAddr, senderMAC: self.devAddr.mac, targetIP: dstIp)
                                await self.routeLayerPacket(dstMac: broadcastMac, type: .arp, data: arpReqeust.marshal())
                                
                                NSLog("[SDLContext] dstIp: \(dstIp) arp query not found")
                            }
                        }
                    }
                }
                
            } while true
        }
        
    }
    
    private func routeLayerPacket(dstMac: Data, type: LayerPacket.PacketType, data: Data) async {
        // 将数据封装层2层的数据包
        let layerPacket = LayerPacket(dstMac: dstMac, srcMac: self.devAddr.mac, type: type, data: data)
        guard let encodedPacket = try? self.aesCipher.encrypt(aesKey: self.aesKey, data: layerPacket.marshal()) else {
            return
        }
        
        // 通过session发送到对端
        if let session = await self.sessionManager.getSession(toAddress: dstMac) {
            NSLog("[SDLContext] send packet by session: \(session)")
            await self.udpHoleActor?.sendPacket(context: self, session: session, data: encodedPacket)
            
            await self.flowTracer.inc(num: data.count, type: .p2p)
        }
        else {
            // 通过super_node进行转发
            await self.udpHoleActor?.forwardPacket(context: self, dst_mac: dstMac, data: encodedPacket)
            // 流量统计
            await self.flowTracer.inc(num: data.count, type: .forward)
            
            // 尝试打洞
            await self.holerManager.addHoler(dstMac: dstMac) {
                self.holerTask(dstMac: dstMac)
            }
        }
    }
    
    func holerTask(dstMac: Data) -> Task<(), Never> {
        return Task {
            guard let message = try? await self.superClientActor?.queryInfo(context: self, dst_mac: dstMac) else {
                return
            }
            
            switch message.packet {
            case .empty:
                SDLLogger.log("[SDLContext] hole query_info get empty: \(message)", level: .debug)
            case .peerInfo(let peerInfo):
                if let remoteAddress = peerInfo.v4Info.socketAddress() {
                    SDLLogger.log("[SDLContext] hole sock address: \(remoteAddress)", level: .warning)
                    // 发送register包
                    await self.udpHoleActor?.sendRegister(context: self, remoteAddress: remoteAddress, dst_mac: dstMac)
                } else {
                    SDLLogger.log("[SDLContext] hole sock address is invalid: \(peerInfo.v4Info)", level: .warning)
                }
            default:
                SDLLogger.log("[SDLContext] hole query_info is packet: \(message)", level: .warning)
            }
        }
    }
    
    deinit {
        self.stunCancel?.cancel()
        self.udpHoleActor = nil
        self.superClientActor = nil
    }
    
}

//--MARK: 获取设备的UUID

extension SDLContext {
    
    public static func getUUID() -> String {
        let userDefaults = UserDefaults.standard
        if let uuid = userDefaults.value(forKey: "gClientId") as? String {
            return uuid
        } else {
            let uuid = UUID().uuidString.replacingOccurrences(of: "-", with: "").lowercased()
            userDefaults.setValue(uuid, forKey: "gClientId")
           
            return uuid
        }
    }
    
    // 获取mac地址
    public static func getMacAddress() -> Data {
        let key = "gMacAddress2"
        
        let userDefaults = UserDefaults.standard
        if let mac = userDefaults.value(forKey: key) as? Data {
            return mac
        }
        else {
            let mac = generateMacAddress()
            userDefaults.setValue(mac, forKey: key)
            
            return mac
        }
    }
    
    // 随机生成mac地址
    private static func generateMacAddress() -> Data {
        var macAddress = [UInt8](repeating: 0, count: 6)
        for i in 0..<6 {
            macAddress[i] = UInt8.random(in: 0...255)
        }
        
        return Data(macAddress)
    }
    
    // 将mac地址转换成字符串
    private static func formatMacAddress(mac: Data) -> String {
        let bytes = [UInt8](mac)
        
        return bytes.map { String(format: "%02X", $0) }.joined(separator: ":").lowercased()
    }
    
}