//
//  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 superService: SDLSuperService
    private let udpHoleService: SDLUDPHoleService
    private let udpHoleV6Service: SDLUDPHoleV6Service
    private let dnsCloudService: DNSCloudService
    private let dnsLocalService: DNSLocalService
    private let flowTracer: SDLFlowTracer

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

    init(provider: NEPacketTunnelProvider,
         config: SDLConfiguration,
         dataCipher: CCDataCipher?,
         sessionManager: SessionManager,
         arpResolver: ArpResolver,
         puncherActor: SDLPuncherActor,
         policyService: PolicyService,
         superService: SDLSuperService,
         udpHoleService: SDLUDPHoleService,
         udpHoleV6Service: SDLUDPHoleV6Service,
         dnsCloudService: DNSCloudService,
         dnsLocalService: DNSLocalService,
         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.superService = superService
        self.udpHoleService = udpHoleService
        self.udpHoleV6Service = udpHoleV6Service
        self.dnsCloudService = dnsCloudService
        self.dnsLocalService = dnsLocalService
        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 runPacketReader() async throws {
        let provider = self.provider

        while !Task.isCancelled {
            guard let batch = await Self.readPackets(from: provider) else {
                break
            }

            try Task.checkCancellation()

            for (data, number) in zip(batch.packets, batch.protocols) where number.int32Value == 2 {
                try Task.checkCancellation()

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

        SDLLogger.log("[PacketOutboundActor] packet reader task finished")
    }

    func handleTunPacket(_ packet: IPPacketView) async {
        let router = PacketOutboundRouter(networkAddress: self.networkAddress, exitNode: self.exitNode)
        let decision = router.route(packet: packet)
        
        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.dnsCloudService.forward(ipPacketData: ipPacketData)
        case .localDNS(let name, let payload, let tracker):
            SDLLogger.log("[PacketOutboundActor] get local dns request: \(name)")
            await self.dnsLocalService.query(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, originalPacket: packet)
        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,
                                        originalPacket: IPPacketView) 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)
            let didSend = await self.routeLayerPacket(dstMac: dstMac, type: type, data: data)
            if didSend {
                self.policyService.recordOutboundFlow(ipPacket: originalPacket)
            }
        } 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.superService.send(type: .arpRequest, data: arpRequest)
            }
        }
    }
    
    @discardableResult
    func routeLayerPacket(dstMac: Data, type: LayerPacket.PacketType, data: Data) async -> Bool {
        guard let plan = try? self.makeDeliveryPlan(dstMac: dstMac, type: type, data: data) else {
            return false
        }

        switch plan {
        case .superNode(let payload):
            await self.sendSuperPacket(type: .data, data: payload)
        case .peer(let payload, let session):
            SDLLogger.log("[PacketOutboundActor] 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] 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.superService.send(type: .queryInfo, data: queryData)
            }
        }

        return true
    }

    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 {
        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("[PacketOutboundActor] unsupported socket family: \(remoteAddress)", for: .debug)
        }
    }
}