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

import Foundation
import NetworkExtension
import NIOCore

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

func startMonitorTask(name: String, _ body: @escaping () async throws -> Void, retryDelay: Duration = .seconds(5)) -> Task<Void, Never> {
    return Task(name: name) {
        while true {
            do {
                try Task.checkCancellation()
                try await body()
            } catch is CancellationError {
                SDLLogger.log("[SDLContext] worker \(name) cancelled", for: .debug)
                break
            } catch let err {
                SDLLogger.log("[SDLContext] worker \(name) crashed: \(err.localizedDescription), will restart", for: .debug)
                do {
                    try await Task.sleep(for: retryDelay)
                } catch is CancellationError {
                    break
                } catch {
                    break
                }
            }
        }
    }
}

enum SDLContextError: Error {
    case udpHoleClosed
    
    case dnsLocalClientClosed
    case dnsLocalClientCancelled
    
    case dnsClientClosed
    case dnsClientCancelled
}

actor SDLContextActor {

    private var config: SDLConfiguration
    // nat的网络类型
    var natType: SDLNATProberActor.NatType = .blocked
    
    // AES加密，授权通过后，对象才会被创建
    private var dataCipher: CCDataCipher?
    
    // session token
    private var sessionToken: Data?
    
    // rsa的相关配置, public_key是本地生成的
    // 加密算法相关
    nonisolated let rsaCipher: RSACipher
    
    private var dnsService: DNSService?
    private let superService: SDLSuperService
    private let udpHoleService: SDLUDPHoleService
    private let packetOutboundActor: PacketOutboundActor
    private let packetInboundActor: PacketInboundActor
    private let tunNetworkManager: SDLTunNetworkManager

    private let publicDnsServers = ["223.5.5.5", "119.29.29.29"]

    nonisolated private let puncherActor: SDLPuncherActor
    // 网络探测对象
    nonisolated private let proberActor: SDLNATProberActor
    
    // 本地ipv6地址信息探测
    private var ipv6AssistClient: SDLIPV6AssistClient?

    private let sessionManager: SessionManager
    nonisolated private let arpResolver: ArpResolver
    
    // 内部socket通讯
    // 改为基于 App Group + Darwin Notification 的通知
    
    // 流量统计
    nonisolated private let flowTracer: SDLFlowTracer
    
    nonisolated private let provider: NEPacketTunnelProvider
    
    // 处理权限控制
    private let policyService: PolicyService
    private var updatePolicyWorker: PeriodicWorker?
    
    // stunRequest任务
    private var stunRequestWorker: PeriodicWorker?
    private var rootTask: Task<Void, Error>?
    private let readySignal = AsyncOneShot<Void>()

    public init(provider: NEPacketTunnelProvider, config: SDLConfiguration, rsaCipher: RSACipher) {
        let puncherActor = SDLPuncherActor()
        let proberActor = SDLNATProberActor(addressArray: config.stunProbeSocketAddressArray)
        let sessionManager = SessionManager()
        let arpResolver = ArpResolver()
        let flowTracer = SDLFlowTracer()
        let policyService = PolicyService(identityId: config.identityId, acl: config.acl)
        let superService = SDLSuperService(serverEndpoint: config.serverEndpoint)
        let udpHoleService = SDLUDPHoleService(proberActor: proberActor)
        let tunNetworkManager = SDLTunNetworkManager(provider: provider)
        let packetOutboundActor = PacketOutboundActor(
            provider: provider,
            config: config,
            dataCipher: nil,
            sessionManager: sessionManager,
            arpResolver: arpResolver,
            puncherActor: puncherActor,
            policyService: policyService,
            superService: superService,
            udpHoleService: udpHoleService,
            flowTracer: flowTracer
        )
        let packetInboundActor = PacketInboundActor(
            provider: provider,
            config: config,
            dataCipher: nil,
            policyService: policyService,
            packetOutboundActor: packetOutboundActor,
            arpResolver: arpResolver,
            superService: superService,
            flowTracer: flowTracer
        )

        self.provider = provider
        self.config = config
        self.rsaCipher = rsaCipher
        
        self.puncherActor = puncherActor
        self.proberActor = proberActor

        self.sessionManager = sessionManager
        self.arpResolver = arpResolver
        self.flowTracer = flowTracer

        // 权限控制
        self.policyService = policyService
        self.superService = superService
        self.udpHoleService = udpHoleService
        self.packetOutboundActor = packetOutboundActor
        self.packetInboundActor = packetInboundActor
        self.tunNetworkManager = tunNetworkManager
    }
    
    public func start() async throws {
        guard self.rootTask == nil else {
            try await self.readySignal.wait(timeout: .seconds(30))
            return
        }

        let rootTask = Task {
            try await self.runRoot()
        }
        self.rootTask = rootTask

        do {
            try await self.readySignal.wait(timeout: .seconds(30))
        } catch {
            rootTask.cancel()
            _ = try? await rootTask.value
            self.rootTask = nil
            throw error
        }
    }

    // 处理context的停止问题
    public func stop() async {
        let rootTask = self.rootTask
        self.rootTask = nil

        rootTask?.cancel()
        _ = try? await rootTask?.value

        await self.cleanupRoot()
    }

    private func runRoot() async throws {
        do {
            try await self.runRootBody()
            await self.cleanupRoot()
        } catch is CancellationError {
            await self.cleanupRoot()
            throw CancellationError()
        } catch {
            await self.readySignal.fail(error)
            await self.cleanupRoot()
            throw error
        }
    }

    private func runRootBody() async throws {
        self.prepareTunnelNotifier()

        // 启动arp的定时清理任务
        await self.puncherActor.start()
        await self.arpResolver.start()

        let dnsService = DNSService(serverIP: self.config.serverEndpoint.ip, publicDnsServers: self.publicDnsServers) { [weak self] event in
            await self?.handleDNSEvent(event)
        }
        self.dnsService = dnsService
        await self.packetOutboundActor.updateDNSService(dnsService)
        await dnsService.start()

        await self.superService.updateMessageHandler { [weak self] message in
            await self?.handleSuperMessage(message: message)
        }

        let packetInboundActor = self.packetInboundActor
        await self.udpHoleService.updateHandlers(
            onEvent: { [weak self] event in
                await self?.handleUDPHoleControlEvent(event)
            },
            onData: { data in
                await packetInboundActor.handleData(data)
            }
        )

        let superService = self.superService
        let udpHoleService = self.udpHoleService
        let packetOutboundActor = self.packetOutboundActor
        let readySignal = self.readySignal

        try await withThrowingTaskGroup(of: Void.self) { group in
            defer {
                group.cancelAll()
            }

            group.addTask {
                try await Self.runRestarting(name: "superService") {
                    try await superService.run()
                }
            }

            group.addTask {
                try await Self.runRestarting(name: "udpHoleService") {
                    try await udpHoleService.run(includeV6: false)
                }
            }

            group.addTask(priority: .high) {
                _ = try await readySignal.wait()
                try await packetOutboundActor.runPacketReader()
            }

            try await group.waitForAll()
        }
    }

    private static func runRestarting(
        name: String,
        retryDelay: Duration = .seconds(5),
        operation: @escaping @Sendable () async throws -> Void
    ) async throws {
        while !Task.isCancelled {
            do {
                try Task.checkCancellation()
                try await operation()
                SDLLogger.log("[SDLContext] worker \(name) ended, will restart", for: .debug)
            } catch is CancellationError {
                SDLLogger.log("[SDLContext] worker \(name) cancelled", for: .debug)
                throw CancellationError()
            } catch {
                SDLLogger.log("[SDLContext] worker \(name) crashed: \(error.localizedDescription), will restart", for: .debug)
            }

            try await Task.sleep(for: retryDelay)
        }
    }

    private func cleanupRoot() async {
        await self.puncherActor.stop()
        await self.arpResolver.stop()
        await self.sessionManager.clear()

        await self.stunRequestWorker?.stop()
        self.stunRequestWorker = nil

        await self.updatePolicyWorker?.stop()
        self.updatePolicyWorker = nil
        await self.policyService.clear()

        await self.udpHoleService.stop()

        let dnsService = self.dnsService
        self.dnsService = nil
        await self.packetOutboundActor.updateDNSService(nil)
        await dnsService?.stop()

        await self.superService.stop()

        self.sessionToken = nil
        self.dataCipher = nil
        self.natType = .blocked
        await self.packetOutboundActor.updateRuntime(config: self.config, dataCipher: nil)
        await self.packetInboundActor.updateRuntime(config: self.config, dataCipher: nil)

        await self.ipv6AssistClient?.stop()
        self.ipv6AssistClient = nil
    }

    deinit {
        SDLLogger.log("[SDLContext] deinit", for: .debug)
    }
    
}

// MARK: probe网络类型
extension SDLContextActor {
    
    private func setNatType(natType: SDLNATProberActor.NatType) {
        self.natType = natType
    }
}

// MARK: Notifier通知机制
extension SDLContextActor {
    
    private func prepareTunnelNotifier() {
        // 启动noticeClient
        // 旧的 UDP NoticeClient 已移除，改为初始化基于 App Group 的通知通道。
        SDLTunnelAppNotifier.shared.clear()
        
        SDLLogger.log("[SDLContext] tunnelAppNotifier ready")
    }
    
    private func publishTunnelEvent(code: Int? = nil, message: String) {
        SDLTunnelAppNotifier.shared.publish(code: code, message: message)
    }
}

// MARK: 数据发送
extension SDLContextActor {
    
    // 发送给super/stun节点的数据
    private func sendSuperPacket(type: SDLPacketType, data: Data) async {
        await self.sendPacket(type: type, data: data, remoteAddress: self.config.stunSocketAddress)
    }
    
    // 发送给peer的数据
    private func sendPeerPacket(type: SDLPacketType, data: Data, remoteAddress: SocketAddress) async {
        await self.sendPacket(type: type, data: data, remoteAddress: remoteAddress)
    }
    
    private func sendPacket(type: SDLPacketType, data: Data, remoteAddress: SocketAddress) async {
        await self.udpHoleService.send(type: type, data: data, remoteAddress: remoteAddress)
    }
}

// MARK: 处理和Super之间的通讯
extension SDLContextActor {
    
    private func handleSuperMessage(message: SDLQUICInboundMessage) async {
        switch message {
        case .welcome(let welcome):
            SDLLogger.log("[SDLContext] quic welcome: \(welcome)")
            
            await self.ipv6AssistClient?.stop()
            self.ipv6AssistClient = nil
            // 尝试创建v6地址的辅助探测器
            if welcome.hasIpv6Assist {
                self.ipv6AssistClient = SDLIPV6AssistClient(assistServerInfo: welcome.ipv6Assist)
                await self.ipv6AssistClient?.start()
            }
            
            // 注册
            await self.doRegisterSuper()
            SDLLogger.log("[SDLContext] quic doRegisterSuper")
        case .pong:
            //SDLLogger.shared.log("[SDLContext] quic pong")
            ()
        case .registerSuperAck(let registerSuperAck):
            await self.handleRegisterSuperAck(registerSuperAck: registerSuperAck)
        case .registerSuperNak(let registerSuperNak):
            await self.handleRegisterSuperNak(nakPacket: registerSuperNak)
        case .peerInfo(let peerInfo):
            SDLLogger.log("[SDLContext] peer message: \(peerInfo)")
            let packets = await self.puncherActor.makeRegisterPackets(peerInfo: peerInfo)
            for packet in packets {
                await self.udpHoleService.send(type: .register, data: packet.data, remoteAddress: packet.remoteAddress)
            }
        case .event(let event):
            await self.handleEvent(event: event)
        case .policyReponse(let policyResponse):
            // 处理权限的请求问题
            await self.policyService.applyPolicyResponse(policyResponse)
        case .exposedServiceResponse(let response):
            await self.applyExposedServiceResponse(response)
        case .arpResponse(let arpResponse):
            SDLLogger.log("[SDLContext] get arp response: \(arpResponse)")
            await self.arpResolver.handleArpResponse(arpResponse: arpResponse)
        }
    }
    
    private func makeSuperEventProcessor() -> SDLSuperEventProcessor {
        return .init(networkAddress: self.config.networkAddress)
    }
    
    private func handleRegisterSuperAck(registerSuperAck: SDLRegisterSuperAck) async {
        // 需要对数据通过rsa的私钥解码
        guard let key = try? self.rsaCipher.decode(data: Data(registerSuperAck.key)) else {
            SDLLogger.log("[SDLContext] registerSuperAck invalid key")
            let error = SDLError.invalidKey
            await self.readySignal.fail(error)
            self.provider.cancelTunnelWithError(error)
            return
        }
        
        let algorithm = registerSuperAck.algorithm.lowercased()
        let regionId = registerSuperAck.regionID
        self.sessionToken = registerSuperAck.sessionToken
        
        switch algorithm {
        case "aes":
            self.dataCipher = CCAESChiper(key: key)
        case "chacha20":
            self.dataCipher = CCChaCha20Cipher(regionId: regionId, keyData: key)
        default:
            SDLLogger.log("[SDLContext] registerSuperAck invalid algorithm \(algorithm)")
            let error = SDLError.unsupportedAlgorithm(algorithm: algorithm)
            await self.readySignal.fail(error)
            self.provider.cancelTunnelWithError(error)
            return
        }
        
        await self.packetOutboundActor.updateRuntime(config: self.config, dataCipher: self.dataCipher)
        await self.packetInboundActor.updateRuntime(config: self.config, dataCipher: self.dataCipher)
        SDLLogger.log("[SDLContext] registerSuperAck, use algorithm \(algorithm), key len: \(key.count)")
        // 服务器分配的tun网卡信息
        do {
            try await self.tunNetworkManager.apply(settings: .init(config: self.config), dnsServer: DNSHelper.dnsServer)
            SDLLogger.log("[SDLContext] setNetworkSettings successed")
            // 开启权限的定时更新
            await self.whenRegistedSuper()
            await self.readySignal.succeed(())
        } catch let err {
            SDLLogger.log("[SDLContext] setTunnelNetworkSettings get error: \(err)")
            await self.readySignal.fail(err)
            self.provider.cancelTunnelWithError(err)
        }
    }
    
    // 注册成功super的回调函数
    private func whenRegistedSuper() async {
        await self.updatePolicyWorker?.stop()
        let policyService = self.policyService
        let superService = self.superService

        let updatePolicyWorker = PeriodicWorker(
            configuration: .init(
                interval: .seconds(10),
                runImmediately: true,
                mode: .fixedDelay,
                errorPolicy: .keepRunning(delay: .seconds(5))
            ),
            operation: {
                SDLLogger.log("[SDLContext] updatePolicyTask execute")
                await policyService.updatePolicy(superService: superService)
            },
            onError: { err in
                SDLLogger.log("[SDLContext] updatePolicyTask stop with err: \(err)")
            }
        )
        self.updatePolicyWorker = updatePolicyWorker
        await updatePolicyWorker.start()
        
        // 启动stun任务
        await self.startStunRequestTask()
    }
    
    private func handleRegisterSuperNak(nakPacket: SDLRegisterSuperNak) async {
        let errorMessage = nakPacket.errorMessage
        guard let errorCode = SDLNAKErrorCode(rawValue: UInt8(nakPacket.errorCode)) else {
            return
        }
        
        switch errorCode {
        case .invalidToken, .nodeDisabled:
            self.publishTunnelEvent(code: Int(errorCode.rawValue), message: errorMessage)
            // 报告错误并退出
            let error = NSError(domain: "com.jihe.punchnet.tun", code: -1)
            await self.readySignal.fail(error)
            self.provider.cancelTunnelWithError(error)
            
        case .noIpAddress, .networkFault, .internalFault:
            self.publishTunnelEvent(code: Int(errorCode.rawValue), message: errorMessage)
        }
        SDLLogger.log("[SDLContext] Get a SuperNak message exit")
    }
    
    private func handleEvent(event: SDLEvent) async {
        let processor = self.makeSuperEventProcessor()
        let plan = await processor.makeProcessingPlan(event: event)
        if let logMessage = plan.logMessage {
            SDLLogger.log(logMessage)
        }
        
        switch plan.action {
        case .removeSession(let dstMac):
            await self.sessionManager.removeSession(dstMac: dstMac)
        case .sendRegister(let registerData, let remoteAddresses):
            for remoteAddress in remoteAddresses {
                await self.sendPeerPacket(type: .register, data: registerData, remoteAddress: remoteAddress)
            }
        case .requestExposedService:
            await self.requestExposedService()
        case .shutdown(let message):
            self.publishTunnelEvent(message: message)
            
            // 报告错误并退出
            let error = NSError(domain: "com.jihe.punchnet.tun", code: -2)
            self.provider.cancelTunnelWithError(error)
        case .none:
            ()
        }
    }
    
    private func doRegisterSuper() async {
        // 注册
        var registerSuper = SDLRegisterSuper()
        registerSuper.clientID = self.config.clientId
        registerSuper.networkID = self.config.networkAddress.networkId
        registerSuper.mac = self.config.networkAddress.mac
        registerSuper.ip = self.config.networkAddress.ip
        registerSuper.maskLen = UInt32(self.config.networkAddress.maskLen)
        registerSuper.hostname = self.config.hostname
        registerSuper.pubKey = self.rsaCipher.pubKey
        registerSuper.accessToken = self.config.accessToken
        
        if let registerSuperData = try? registerSuper.serializedData() {
            SDLLogger.log("[SDLContext] will send register super")
            await self.superService.send(type: .registerSuper, data: registerSuperData)
        }
    }

    private func requestExposedService() async {
        guard let requestData = await self.policyService.makeExposedServiceRequest() else {
            return
        }

        await self.superService.send(type: .exposedServiceRequest, data: requestData)
    }

    private func applyExposedServiceResponse(_ response: SDLExposedServiceResponse) async {
        guard let acl = await self.policyService.applyExposedServiceResponse(response) else {
            return
        }

        self.config.acl = acl
    }
    
}

// MARK: DNS service events
extension SDLContextActor {
    private func handleDNSEvent(_ event: DNSService.Event) async {
        switch event {
        case .packet(let packet):
            let nePacket = NEPacket(data: packet, protocolFamily: 2)
            self.provider.packetFlow.writePacketObjects([nePacket])
        }
    }
}

// MARK: 处理从Hole收到的数据
extension SDLContextActor {
    private func handleUDPHoleControlEvent(_ event: SDLUDPHoleService.Event) async {
        switch event {
        case .ready(let localAddress):
            SDLLogger.log("[SDLContext] udpHole ready: \(localAddress)")
        case .natType(let natType):
            self.setNatType(natType: natType)
            SDLLogger.log("[SDLContext] nat_type is: \(natType)")
        case .packet(let remoteAddress, let message, let source):
            await self.handleUDPHolePacket(remoteAddress: remoteAddress, message: message, source: source)
        case .closed(let error):
            SDLLogger.log("[SDLContext] udpHole closed: \(error)", for: .debug)
        }
    }

    private func handleUDPHolePacket(remoteAddress: SocketAddress, message: SDLHoleControlMessage, source: SDLUDPHoleKind) async {
        switch message {
        case .stunReply(_), .stunProbeReply(_):
            SDLLogger.log("[SDLContext] get a stun reply", for: .debug)
        case .register(let register):
            try? await self.handleRegister(remoteAddress: remoteAddress, register: register, source: source)
        case .registerAck(let registerAck):
            await self.handleRegisterAck(remoteAddress: remoteAddress, registerAck: registerAck, source: source)
        }
    }
    
    private func handleRegister(remoteAddress: SocketAddress, register: SDLRegister, source: SDLUDPHoleKind) async throws {
        let networkAddr = config.networkAddress
        SDLLogger.log("[SDLContext] register packet: \(register), network_address: \(networkAddr)")
        
        // 判断目标地址是否是tun的网卡地址, 并且是在同一个网络下
        if register.dstMac == networkAddr.mac && register.networkID == networkAddr.networkId {
            // 回复ack包
            var registerAck = SDLRegisterAck()
            registerAck.networkID = networkAddr.networkId
            registerAck.srcMac = networkAddr.mac
            registerAck.dstMac = register.srcMac
            
            await self.sendPeerPacket(type: .registerAck, data: try registerAck.serializedData(), remoteAddress: remoteAddress)
            // 这里需要建立到来源的会话, 在复杂网络下，通过super-node查询到的nat地址不一定靠谱，需要通过udp包的来源地址作为nat地址
            if let session = Session(dstMac: register.srcMac, natAddress: remoteAddress, addressType: source.convertAddressType()) {
                await self.sessionManager.addSession(session: session)
            } else {
                SDLLogger.log("[SDLContext] didReadRegister get unsupported remoteAddress: \(remoteAddress)", for: .debug)
            }
        } else {
            SDLLogger.log("[SDLContext] didReadRegister get a invalid packet, because dst_ip not matched: \(register.dstMac)")
        }
    }
    
    private func handleRegisterAck(remoteAddress: SocketAddress, registerAck: SDLRegisterAck, source: SDLUDPHoleKind) async {
        // 判断目标地址是否是tun的网卡地址, 并且是在同一个网络下
        let networkAddr = config.networkAddress
        if registerAck.dstMac == networkAddr.mac && registerAck.networkID == networkAddr.networkId {
            if let session = Session(dstMac: registerAck.srcMac, natAddress: remoteAddress, addressType: source.convertAddressType()) {
                await self.sessionManager.addSession(session: session)
            } else {
                SDLLogger.log("[SDLContext] didReadRegisterAck get unsupported remoteAddress: \(remoteAddress)", for: .debug)
            }
        } else {
            SDLLogger.log("[SDLContext] didReadRegisterAck get a invalid packet, because dst_mac not matched: \(registerAck.dstMac)")
        }
    }
    
}

// MARK: 和Stun相关的心跳机制
extension SDLContextActor {
    
    // MARK: -- StunRequestTask
    private func startStunRequestTask() async {
        await self.stunRequestWorker?.stop()

        let stunRequestWorker = PeriodicWorker(
            configuration: .init(
                interval: .seconds(8),
                runImmediately: true,
                mode: .fixedDelay,
                errorPolicy: .keepRunning(delay: .seconds(5))
            ),
            operation: { [weak self] in
                let probeReply = try? await self?.ipv6AssistClient?.probe(requestTimeout: .seconds(3))

                if let v6Info = probeReply?.v6Info, let v6Address = SDLUtil.ipv6DataToString(v6Info.v6) {
                    SDLLogger.log("[SDLContext] probe ipv6 address: \(v6Address)")
                } else {
                    SDLLogger.log("[SDLContext] probe ipv6 address: empty")
                }

                await self?.sendStunRequest(v6Info: probeReply?.v6Info)
            },
            onError: { err in
                SDLLogger.log("[SDLContext] udp stunRequestTask stop with err: \(err)")
            }
        )
        self.stunRequestWorker = stunRequestWorker
        await stunRequestWorker.start()
    }
    
    private func sendStunRequest(v6Info: SDLV6Info?) async {
        guard let sessionToken = self.sessionToken else {
            return
        }
        
        var stunRequest = SDLStunRequest()
        stunRequest.clientID = self.config.clientId
        stunRequest.networkID = self.config.networkAddress.networkId
        stunRequest.ip = self.config.networkAddress.ip
        stunRequest.mac = self.config.networkAddress.mac
        stunRequest.natType = UInt32(self.natType.rawValue)
        stunRequest.sessionToken = sessionToken
        
        if let v6Info {
            stunRequest.v6Info = v6Info
        }
        
        if let stunData = try? stunRequest.serializedData() {
            await self.sendSuperPacket(type: .stunRequest, data: stunData)
        }
    }
}

// MARK: NEPacketTunnelProvider相关的逻辑
extension SDLContextActor {
    
    // 取消出口节点的时候，ip地址为: 0.0.0.0
    public func updateExitNode(exitNodeIp: String) async throws {
        if let ip = SDLUtil.ipv4StrToInt32(exitNodeIp), ip > 0 {
            self.config.exitNode = .init(exitNodeIp: ip)
        } else {
            self.config.exitNode = nil
        }
        await self.packetOutboundActor.updateRuntime(config: self.config, dataCipher: self.dataCipher)
        await self.packetInboundActor.updateRuntime(config: self.config, dataCipher: self.dataCipher)
        
        try await self.tunNetworkManager.apply(settings: .init(config: self.config), dnsServer: DNSHelper.dnsServer)
    }
    
}