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

import Foundation
import NetworkExtension
import NIOCore

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

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

    private static 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 ipv6AssistEvents: AsyncStream<SDLV6Info?>
    private let ipv6AssistContinuation: AsyncStream<SDLV6Info?>.Continuation

    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 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 udpHoleV6Service = SDLUDPHoleV6Service()
        let dnsCloudService = DNSCloudService(serverIP: config.serverEndpoint.ip)
        let dnsLocalService = DNSLocalService(publicDnsServers: Self.publicDnsServers)
        let superControlPlane = SDLSuperControlPlane(config: config, rsaCipher: rsaCipher)
        let tunNetworkManager = SDLTunNetworkManager(provider: provider)
        let ipv6AssistPair = AsyncStream.makeStream(of: Optional<SDLV6Info>.self, bufferingPolicy: .bufferingNewest(1))
        let packetOutboundActor = PacketOutboundActor(
            provider: provider,
            config: config,
            dataCipher: nil,
            sessionManager: sessionManager,
            arpResolver: arpResolver,
            puncherActor: puncherActor,
            policyService: policyService,
            superService: superService,
            udpHoleService: udpHoleService,
            udpHoleV6Service: udpHoleV6Service,
            dnsCloudService: dnsCloudService,
            dnsLocalService: dnsLocalService,
            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.dnsCloudService = dnsCloudService
        self.dnsLocalService = dnsLocalService
        self.superService = superService
        self.superControlPlane = superControlPlane
        self.udpHoleService = udpHoleService
        self.udpHoleV6Service = udpHoleV6Service
        self.packetOutboundActor = packetOutboundActor
        self.packetInboundActor = packetInboundActor
        self.tunNetworkManager = tunNetworkManager
        self.ipv6AssistEvents = ipv6AssistPair.stream
        self.ipv6AssistContinuation = ipv6AssistPair.continuation
    }
    
    public func start() async throws {
        guard self.rootTask == nil else {
            try await self.readySignal.wait(timeout: .seconds(30))
            return
        }

        let rootTask = Task {
            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
            }
        }
        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 runRootBody() async throws {
        self.prepareTunnelNotifier()

        await self.dnsCloudService.updateEventHandler { [weak self] event in
            await self?.handleDNSEvent(event)
        }
        
        await self.dnsLocalService.updateEventHandler { [weak self] event in
            await self?.handleDNSEvent(event)
        }

        await self.superControlPlane.updateDecisionHandler { [weak self] decision in
            await self?.handleSuperDecision(decision)
        }

        let superControlPlane = self.superControlPlane
        await self.superService.updateMessageHandler { message in
            await superControlPlane.handle(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)
            }
        )
        
        await self.udpHoleV6Service.updateHandlers(
            onEvent: { [weak self] event in
                await self?.handleUDPHoleControlEvent(event)
            },
            onData: { data in
                await packetInboundActor.handleData(data)
            }
        )

        let superService = self.superService
        let dnsCloudService = self.dnsCloudService
        let dnsLocalService = self.dnsLocalService
        let udpHoleService = self.udpHoleService
        let udpHoleV6Service = self.udpHoleV6Service
        let packetOutboundActor = self.packetOutboundActor
        let policyService = self.policyService
        let puncherActor = self.puncherActor
        let arpResolver = self.arpResolver
        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()
                }
            }

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

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

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

            group.addTask {
                try await puncherActor.runCleanup()
            }

            group.addTask {
                try await arpResolver.runCleanup()
            }

            group.addTask {
                try await self.runIPv6AssistSupervisor()
            }

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

            group.addTask {
                _ = try await readySignal.wait()
                try await Self.runPeriodic(name: "updatePolicyTask", interval: .seconds(10)) {
                    SDLLogger.log("[SDLContext] updatePolicyTask execute")
                    await policyService.updatePolicy(superService: superService)
                }
            }

            group.addTask {
                _ = try await readySignal.wait()
                try await Self.runPeriodic(name: "stunRequestTask", interval: .seconds(8)) {
                    try await self.runStunRequestOnce()
                }
            }

            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 static func runPeriodic(
        name: String,
        interval: Duration,
        retryDelay: Duration = .seconds(5),
        operation: @escaping @Sendable () async throws -> Void
    ) async throws {
        while !Task.isCancelled {
            do {
                try Task.checkCancellation()
                try await operation()
                try await Task.sleep(for: interval)
            } catch is CancellationError {
                SDLLogger.log("[SDLContext] worker \(name) cancelled", for: .debug)
                throw CancellationError()
            } catch {
                SDLLogger.log("[SDLContext] worker \(name) crashed: \(error.localizedDescription), will retry", for: .debug)
                try await Task.sleep(for: retryDelay)
            }
        }
    }

    private func runIPv6AssistSupervisor() async throws {
        do {
            for await assistInfo in self.ipv6AssistEvents {
                try Task.checkCancellation()
                await self.stopCurrentIPv6AssistClient()

                guard let assistInfo else {
                    continue
                }

                guard let client = SDLIPV6AssistClient(assistServerInfo: assistInfo) else {
                    SDLLogger.log("[SDLContext] invalid ipv6 assist config", for: .debug)
                    continue
                }

                self.ipv6AssistClient = client

                do {
                    try await client.run()
                } catch is CancellationError {
                    throw CancellationError()
                } catch {
                    SDLLogger.log("[SDLContext] ipv6 assist client ended: \(error.localizedDescription)", for: .debug)
                }

                if self.ipv6AssistClient === client {
                    self.ipv6AssistClient = nil
                }
            }
        } catch is CancellationError {
            await self.stopCurrentIPv6AssistClient()
            throw CancellationError()
        }

        await self.stopCurrentIPv6AssistClient()
    }

    private func stopCurrentIPv6AssistClient() async {
        let client = self.ipv6AssistClient
        self.ipv6AssistClient = nil
        await client?.stop()
    }

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

        await self.policyService.clear()

        await self.udpHoleService.stop()
        await self.udpHoleV6Service.stop()

        await self.dnsCloudService.stop()
        await self.dnsLocalService.stop()

        await self.superService.stop()
        await self.superControlPlane.reset()

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

        self.ipv6AssistContinuation.yield(nil)
        await self.stopCurrentIPv6AssistClient()
    }

    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 {
    
    // 发送给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 {
        switch remoteAddress {
        case .v4:
            await self.udpHoleService.send(type: type, data: data, remoteAddress: remoteAddress)
        case .v6:
            await self.udpHoleV6Service.send(type: type, data: data, remoteAddress: remoteAddress)
        default:
            SDLLogger.log("[SDLContext] unsupported socket family: \(remoteAddress)", for: .debug)
        }
    }
}

// MARK: 处理和Super之间的通讯
extension SDLContextActor {

    private func handleSuperDecision(_ decision: SuperDecision) async {
        switch decision {
        case .updateIPv6Assist(let assistInfo):
            await self.stopCurrentIPv6AssistClient()
            self.ipv6AssistContinuation.yield(assistInfo)
        case .completeRegistration(let cipher):
            await self.completeSuperRegistration(cipher: cipher)
        case .failTunnel(let error):
            await self.failTunnel(error)
        case .publishTunnelEvent(let code, let message):
            self.publishTunnelEvent(code: code, message: message)
        case .sendSuper(let type, let data):
            await self.superService.send(type: type, data: data)
        case .sendPacket(let type, let data, let remoteAddress):
            await self.sendPacket(type: type, data: data, remoteAddress: remoteAddress)
        case .resolvePeerInfo(let peerInfo):
            let packets = await self.puncherActor.makeRegisterPackets(peerInfo: peerInfo)
            for packet in packets {
                await self.sendPacket(type: .register, data: packet.data, remoteAddress: packet.remoteAddress)
            }
        case .removeSession(let dstMac):
            await self.sessionManager.removeSession(dstMac: dstMac)
        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 .applyPolicyResponse(let policyResponse):
            await self.policyService.applyPolicyResponse(policyResponse)
        case .applyExposedServiceResponse(let response):
            await self.applyExposedServiceResponse(response)
        case .handleARPResponse(let arpResponse):
            await self.arpResolver.handleArpResponse(arpResponse: arpResponse)
        }
    }

    private func completeSuperRegistration(cipher: CCDataCipher) async {
        self.dataCipher = cipher
        await self.packetOutboundActor.updateRuntime(config: self.config, dataCipher: cipher)
        await self.packetInboundActor.updateRuntime(config: self.config, dataCipher: cipher)

        do {
            try await self.tunNetworkManager.apply(settings: .init(config: self.config), dnsServer: DNSHelper.dnsServer)
            SDLLogger.log("[SDLContext] setNetworkSettings successed")
            await self.readySignal.succeed(())
        } catch {
            SDLLogger.log("[SDLContext] setTunnelNetworkSettings get error: \(error)")
            await self.failTunnel(error)
        }
    }

    private func failTunnel(_ error: Error) async {
        await self.readySignal.fail(error)
        self.provider.cancelTunnelWithError(error)
    }

    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: DNSEvent) 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):
            await self.handleUDPHolePacket(remoteAddress: remoteAddress, message: message)
        case .closed(let error):
            SDLLogger.log("[SDLContext] udpHole closed: \(error)", for: .debug)
        }
    }

    private func handleUDPHolePacket(remoteAddress: SocketAddress, message: SDLHoleControlMessage) 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)
        case .registerAck(let registerAck):
            await self.handleRegisterAck(remoteAddress: remoteAddress, registerAck: registerAck)
        }
    }
    
    private func handleRegister(remoteAddress: SocketAddress, register: SDLRegister) 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 addressType = Self.addressType(from: remoteAddress),
               let session = Session(dstMac: register.srcMac, natAddress: remoteAddress, addressType: addressType) {
                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) async {
        // 判断目标地址是否是tun的网卡地址, 并且是在同一个网络下
        let networkAddr = config.networkAddress
        if registerAck.dstMac == networkAddr.mac && registerAck.networkID == networkAddr.networkId {
            if let addressType = Self.addressType(from: remoteAddress),
               let session = Session(dstMac: registerAck.srcMac, natAddress: remoteAddress, addressType: addressType) {
                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)")
        }
    }

    private static func addressType(from remoteAddress: SocketAddress) -> Session.AddressType? {
        switch remoteAddress {
        case .v4:
            return .v4
        case .v6:
            return .v6
        default:
            return nil
        }
    }
    
}

// MARK: 和Stun相关的心跳机制
extension SDLContextActor {
    private func runStunRequestOnce() async throws {
        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.superControlPlane.sendStunRequest(natType: self.natType, v6Info: probeReply?.v6Info)
    }
}

// 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)
    }
    
}