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

import Foundation
import NetworkExtension

actor PacketInboundActor {
    private struct PolicyPacketContext {
        let srcIdentityID: UInt32
        let proto: UInt8
        let srcIP: UInt32
        let dstIP: UInt32
        let srcPort: UInt16?
        let dstPort: UInt16?

        var logDescription: String {
            let srcPort = self.srcPort.map(String.init) ?? "-"
            let dstPort = self.dstPort.map(String.init) ?? "-"
            return "srcIdentityID: \(self.srcIdentityID), proto: \(self.proto), srcIP: \(SDLUtil.int32ToIp(self.srcIP)), dstIP: \(SDLUtil.int32ToIp(self.dstIP)), srcPort: \(srcPort), dstPort: \(dstPort)"
        }
    }

    private enum ProcessingAction {
        case sendARPReply(dstMac: Data, data: Data)
        case appendARP(ip: UInt32, mac: Data)
        case writeToTun(packetData: Data, identityID: UInt32)
        case requestPolicy(PolicyPacketContext)
        case dropByPolicy(PolicyPacketContext)
        case none
    }

    private struct ProcessingPlan {
        let inboundBytes: Int
        let action: ProcessingAction
    }

    private let provider: NEPacketTunnelProvider
    private let policyService: PolicyService
    private let packetOutboundActor: PacketOutboundActor
    private let arpResolver: ArpResolver
    private let superService: SDLSuperService
    private let flowTracer: SDLFlowTracer

    private var networkAddress: SDLConfiguration.NetworkAddress
    private var identityId: UInt32
    private var dataCipher: CCDataCipher?

    init(provider: NEPacketTunnelProvider,
         config: SDLConfiguration,
         dataCipher: CCDataCipher?,
         policyService: PolicyService,
         packetOutboundActor: PacketOutboundActor,
         arpResolver: ArpResolver,
         superService: SDLSuperService,
         flowTracer: SDLFlowTracer) {
        self.provider = provider
        self.networkAddress = config.networkAddress
        self.identityId = config.identityId
        self.dataCipher = dataCipher
        self.policyService = policyService
        self.packetOutboundActor = packetOutboundActor
        self.arpResolver = arpResolver
        self.superService = superService
        self.flowTracer = flowTracer
    }

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

    func handleData(_ data: SDLData) async {
        let policyRuntime = self.policyService.policyRuntime()
        guard let plan = try? self.makeProcessingPlan(data: data, policyRuntime: policyRuntime) else {
            return
        }

        self.flowTracer.inc(num: plan.inboundBytes, type: .inbound)

        switch plan.action {
        case .sendARPReply(let dstMac, let responseData):
            SDLLogger.log("[PacketInboundActor] get arp request packet", category: .packet)
            await self.packetOutboundActor.routeLayerPacket(dstMac: dstMac, type: .arp, data: responseData)
        case .appendARP(let ip, let mac):
            SDLLogger.log("[PacketInboundActor] get arp response packet", category: .packet)
            await self.arpResolver.append(ip: ip, mac: mac)
        case .writeToTun(let packetData, let identityID):
            let packet = NEPacket(data: packetData, protocolFamily: 2)
            self.provider.packetFlow.writePacketObjects([packet])
            SDLLogger.trace("[PacketInboundActor] hole identity: \(identityID), allow, data count: \(packetData.count)", category: .packet)
        case .requestPolicy(let context):
            SDLLogger.log("[PacketInboundActor] policy miss, \(context.logDescription)", category: .packet)
            if let queryData = await self.policyService.makePolicyRequest(srcIdentityID: context.srcIdentityID) {
                await self.superService.send(type: .policyRequest, data: queryData)
            }
        case .dropByPolicy(let context):
            SDLLogger.trace("[PacketInboundActor] policy denied, \(context.logDescription)", category: .packet)
        case .none:
            ()
        }
    }

    private func makeProcessingPlan(data: SDLData, policyRuntime: PolicyRuntime) throws -> ProcessingPlan? {
        guard let dataCipher = self.dataCipher else {
            return nil
        }

        let mac = LayerPacket.MacAddress(data: data.dstMac)
        guard (data.dstMac == self.networkAddress.mac || mac.isBroadcast() || mac.isMulticast()) else {
            return nil
        }

        let decryptedData = try dataCipher.decrypt(cipherText: data.data)
        let layerPacket = try LayerPacketView(layerData: decryptedData)
        let inboundBytes = decryptedData.count

        switch layerPacket.type {
        case .arp:
            return self.makeARPPlan(layerData: layerPacket.data, inboundBytes: inboundBytes)
        case .ipv4:
            return self.makeIPv4Plan(
                layerData: layerPacket.data,
                identityID: data.identityID,
                inboundBytes: inboundBytes,
                policyRuntime: policyRuntime
            )
        default:
            SDLLogger.log("[SDLContext] get invalid packet", category: .packet)
            return .init(inboundBytes: inboundBytes, action: .none)
        }
    }

    private func makeARPPlan(layerData: Data, inboundBytes: Int) -> ProcessingPlan {
        if let arpPacket = ARPPacket(data: layerData) {
            if arpPacket.targetIP == self.networkAddress.ip {
                switch arpPacket.opcode {
                case .request:
                    let response = ARPPacket.arpResponse(for: arpPacket, mac: self.networkAddress.mac, ip: self.networkAddress.ip)
                    return .init(
                        inboundBytes: inboundBytes,
                        action: .sendARPReply(dstMac: arpPacket.senderMAC, data: response.marshal())
                    )
                case .response:
                    return .init(
                        inboundBytes: inboundBytes,
                        action: .appendARP(ip: arpPacket.senderIP, mac: arpPacket.senderMAC)
                    )
                }
            } else {
                SDLLogger.log("[SDLContext] get invalid arp packet: \(arpPacket), target_ip: \(SDLUtil.int32ToIp(arpPacket.targetIP)), net ip: \(SDLUtil.int32ToIp(self.networkAddress.ip))", category: .packet)
            }
        } else {
            SDLLogger.log("[SDLContext] get invalid arp packet", category: .packet)
        }

        return .init(inboundBytes: inboundBytes, action: .none)
    }

    private func makeIPv4Plan(layerData: Data, identityID: UInt32, inboundBytes: Int, policyRuntime: PolicyRuntime) -> ProcessingPlan {
        guard let ipPacket = IPPacketView(layerData) else {
            return .init(inboundBytes: inboundBytes, action: .none)
        }

        switch policyRuntime.evaluateInbound(srcIdentityID: identityID, ipPacket: ipPacket) {
        case .allow:
            return .init(
                inboundBytes: inboundBytes,
                action: .writeToTun(packetData: ipPacket.data, identityID: identityID)
            )
        case .deny:
            return .init(
                inboundBytes: inboundBytes,
                action: .dropByPolicy(self.makePolicyPacketContext(identityID: identityID, ipPacket: ipPacket))
            )
        case .missingPolicy:
            return .init(
                inboundBytes: inboundBytes,
                action: .requestPolicy(self.makePolicyPacketContext(identityID: identityID, ipPacket: ipPacket))
            )
        }
    }

    private func makePolicyPacketContext(identityID: UInt32, ipPacket: IPPacketView) -> PolicyPacketContext {
        let ports: (UInt16?, UInt16?)
        switch ipPacket.transportPacket {
        case .tcp(let srcPort, let dstPort, _):
            ports = (srcPort, dstPort)
        case .udp(let srcPort, let dstPort, _):
            ports = (srcPort, dstPort)
        default:
            ports = (nil, nil)
        }

        return PolicyPacketContext(
            srcIdentityID: identityID,
            proto: ipPacket.header.proto,
            srcIP: ipPacket.header.source,
            dstIP: ipPacket.header.destination,
            srcPort: ports.0,
            dstPort: ports.1
        )
    }
}