//
//  PacketOutboundActor.swift
//  Tun
//
//  Created by Codex on 2026/5/20.
//

import Foundation
import NetworkExtension
import NIOCore

actor PacketOutboundActor {
    private typealias PacketReadResult = (packets: [Data], protocols: [NSNumber])?

    private enum DeliveryPlan {
        case superNode(payload: Data)
        case peer(payload: Data, session: Session)
        case superNodeAndPunch(payload: Data, request: SDLPuncherActor.RegisterRequest)
    }

    private let provider: NEPacketTunnelProvider
    private let sessionManager: SessionManager
    private let arpResolver: ArpResolver
    private let puncherActor: SDLPuncherActor
    private let policyService: PolicyService
    private let superServiceProxy: SDLSuperServiceProxy
    private let udpHoleServiceProxy: SDLUDPHoleServiceProxy
    private let flowTracer: SDLFlowTracer
    private var packetReaderTask: Task<Void, Never>?
    private var packetReaderGeneration: UInt64 = 0

    private var networkAddress: SDLConfiguration.NetworkAddress
    private var identityId: UInt32
    private var exitNode: SDLConfiguration.ExitNode?
    private var stunSocketAddress: SocketAddress
    private var dataCipher: CCDataCipher?
    private var dnsService: DNSService?

    init(provider: NEPacketTunnelProvider,
         config: SDLConfiguration,
         dataCipher: CCDataCipher?,
         sessionManager: SessionManager,
         arpResolver: ArpResolver,
         puncherActor: SDLPuncherActor,
         policyService: PolicyService,
         superServiceProxy: SDLSuperServiceProxy,
         udpHoleServiceProxy: SDLUDPHoleServiceProxy,
         flowTracer: SDLFlowTracer) {
        self.provider = provider
        self.networkAddress = config.networkAddress
        self.identityId = config.identityId
        self.exitNode = config.exitNode
        self.stunSocketAddress = config.stunSocketAddress
        self.dataCipher = dataCipher
        self.sessionManager = sessionManager
        self.arpResolver = arpResolver
        self.puncherActor = puncherActor
        self.policyService = policyService
        self.superServiceProxy = superServiceProxy
        self.udpHoleServiceProxy = udpHoleServiceProxy
        self.flowTracer = flowTracer
    }

    func updateRuntime(config: SDLConfiguration, dataCipher: CCDataCipher?) {
        self.networkAddress = config.networkAddress
        self.identityId = config.identityId
        self.exitNode = config.exitNode
        self.stunSocketAddress = config.stunSocketAddress
        self.dataCipher = dataCipher
    }

    func updateDNSService(_ dnsService: DNSService?) {
        self.dnsService = dnsService
    }

    func startPacketReader() async {
        guard self.packetReaderTask == nil else {
            return
        }

        self.packetReaderGeneration &+= 1

        let provider = self.provider
        let generation = self.packetReaderGeneration

        self.packetReaderTask = Task(priority: .high) { [weak self] in
            while !Task.isCancelled {
                guard let batch = await Self.readPackets(from: provider) else {
                    break
                }

                if Task.isCancelled {
                    break
                }

                for (data, number) in zip(batch.packets, batch.protocols) where number.int32Value == 2 {
                    if Task.isCancelled {
                        break
                    }

                    if let packet = IPPacket(data) {
                        await self?.handleTunPacket(packet)
                    }
                }
            }

            await self?.finishPacketReader(generation: generation)
            SDLLogger.log("[PacketOutboundActor] packet reader task finished")
        }
    }

    func stop() async {
        self.packetReaderGeneration &+= 1

        let packetReaderTask = self.packetReaderTask
        self.packetReaderTask = nil
        packetReaderTask?.cancel()
    }

    func handleTunPacket(_ packet: IPPacket) async {
        let router = PacketOutboundRouter(networkAddress: self.networkAddress, exitNode: self.exitNode)
        let decision = router.route(packet: packet)

        if decision.shouldTrackFlow {
            self.policyService.recordOutboundFlow(ipPacket: packet)
        }

        await self.handleTunRouteDecision(decision)
    }
    
    private func handleTunRouteDecision(_ decision: PacketOutboundRouter.RouteDecision) async {
        switch decision {
        case .loopback(let ipPacketData):
            let nePacket = NEPacket(data: ipPacketData, protocolFamily: 2)
            self.provider.packetFlow.writePacketObjects([nePacket])
        case .cloudDNS(let name, let ipPacketData):
            SDLLogger.log("[PacketOutboundActor] get cloud dns request: \(name)")
            await self.dnsService?.forward(ipPacketData: ipPacketData)
        case .localDNS(let name, let payload, let tracker):
            SDLLogger.log("[PacketOutboundActor] get local dns request: \(name)")
            await self.dnsService?.queryLocal(tracker: tracker, dnsPayload: payload)
        case .forwardToNextHop(let ip, let type, let data, let kind):
            await self.forwardPacketToNextHop(ip: ip, type: type, data: data, kind: kind)
        case .drop(let reason):
            SDLLogger.log("[PacketOutboundActor] drop tun packet, reason: \(reason.rawValue)", for: .trace)
        }
    }
    
    private func forwardPacketToNextHop(ip: UInt32, type: LayerPacket.PacketType, data: Data, kind: PacketOutboundRouter.ForwardKind) async {
        switch kind {
        case .sameNetwork:
            SDLLogger.log("[PacketOutboundActor] dstIp: \(SDLUtil.int32ToIp(ip)) same network", for: .trace)
        case .exitNode, .dnsExitNode:
            SDLLogger.log("[PacketOutboundActor] use exit_node: \(SDLUtil.int32ToIp(ip))", for: .trace)
        }

        if let dstMac = self.arpResolver.snapshot().lookup(ip) {
            SDLLogger.log("[PacketOutboundActor] dstIp: \(SDLUtil.int32ToIp(ip)), dst_mac is: \(SDLUtil.formatMacAddress(mac: dstMac))", for: .trace)
            await self.routeLayerPacket(dstMac: dstMac, type: type, data: data)
        } else {
            SDLLogger.log("[PacketOutboundActor] dstIp: \(SDLUtil.int32ToIp(ip)) arp query not found, broadcast", for: .trace)
            if let arpRequest = try? await self.arpResolver.makeArpRequest(targetIp: ip) {
                await self.superServiceProxy.send(type: .arpRequest, data: arpRequest)
            }
        }
    }
    
    func routeLayerPacket(dstMac: Data, type: LayerPacket.PacketType, data: Data) async {
        guard let plan = try? self.makeDeliveryPlan(dstMac: dstMac, type: type, data: data) else {
            return
        }

        switch plan {
        case .superNode(let payload):
            await self.sendSuperPacket(type: .data, data: payload)
        case .peer(let payload, let session):
            SDLLogger.log("[PacketOutboundActor] step 5 send packet by session: \(session)", for: .trace)
            await self.sendPeerPacket(type: .data, data: payload, remoteAddress: session.natAddress)
            self.flowTracer.inc(num: payload.count, type: .p2p)
        case .superNodeAndPunch(let payload, let request):
            await self.sendSuperPacket(type: .data, data: payload)
            SDLLogger.log("[PacketOutboundActor] step 5 send packet by super: \(self.stunSocketAddress)", for: .trace)
            self.flowTracer.inc(num: payload.count, type: .forward)

            if let queryData = await self.puncherActor.makeQueryInfoRequest(request: request) {
                await self.superServiceProxy.send(type: .queryInfo, data: queryData)
            }
        }
    }

    private func finishPacketReader(generation: UInt64) {
        guard generation == self.packetReaderGeneration else {
            return
        }

        self.packetReaderTask = nil
    }

    private static func readPackets(from provider: NEPacketTunnelProvider) async -> PacketReadResult {
        let readContinuation = OnceContinuation<PacketReadResult, Never>()

        return await withTaskCancellationHandler {
            await withCheckedContinuation { continuation in
                readContinuation.set(continuation)
                provider.packetFlow.readPackets { packets, protocols in
                    readContinuation.resume(returning: (packets: packets, protocols: protocols))
                }
            }
        } onCancel: {
            readContinuation.resume(returning: nil)
        }
    }
    
    private func makeDeliveryPlan(dstMac: Data, type: LayerPacket.PacketType, data: Data) throws -> DeliveryPlan? {
        guard let payload = try self.makeDataPayload(dstMac: dstMac, type: type, data: data) else {
            return nil
        }

        if ARPPacket.isBroadcastMac(dstMac) {
            return .superNode(payload: payload)
        }

        if let session = self.sessionManager.snapshot().getSession(toAddress: dstMac) {
            return .peer(payload: payload, session: session)
        }

        return .superNodeAndPunch(
            payload: payload,
            request: .init(
                srcMac: self.networkAddress.mac,
                dstMac: dstMac,
                networkId: self.networkAddress.networkId
            )
        )
    }

    private func makeDataPayload(dstMac: Data, type: LayerPacket.PacketType, data: Data) throws -> Data? {
        guard let dataCipher = self.dataCipher else {
            return nil
        }

        let layerPacket = LayerPacket(dstMac: dstMac, srcMac: self.networkAddress.mac, type: type, data: data)
        let encodedPacket = try dataCipher.encrypt(plainText: layerPacket.marshal())

        var dataPacket = SDLData()
        dataPacket.networkID = self.networkAddress.networkId
        dataPacket.srcMac = self.networkAddress.mac
        dataPacket.dstMac = dstMac
        dataPacket.ttl = 255
        dataPacket.identityID = self.identityId
        dataPacket.isP2P = true
        dataPacket.data = encodedPacket

        return try dataPacket.serializedData()
    }

    private func sendSuperPacket(type: SDLPacketType, data: Data) async {
        await self.sendPacket(type: type, data: data, remoteAddress: self.stunSocketAddress)
    }

    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.udpHoleServiceProxy.send(type: type, data: data, remoteAddress: remoteAddress)
    }
}