//
//  PacketInboundProcessor.swift
//  Tun
//
//  Created by 安礼成 on 2026/4/14.
//

import Foundation

final class PacketInboundProcessor {
    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)"
        }
    }
    
    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
    }
    
    struct ProcessingPlan {
        let inboundBytes: Int
        let action: ProcessingAction
    }
    
    private let networkAddress: SDLConfiguration.NetworkAddress
    private let dataCipher: CCDataCipher?
    private let policyRuntime: PolicyRuntime
    
    init(networkAddress: SDLConfiguration.NetworkAddress,
         dataCipher: CCDataCipher?,
         policyRuntime: PolicyRuntime) {
        self.networkAddress = networkAddress
        self.dataCipher = dataCipher
        self.policyRuntime = policyRuntime
    }
    
    func makeProcessingPlan(data: SDLData) 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.data))
        let layerPacket = try LayerPacket(layerData: decryptedData)
        let inboundBytes = decryptedData.count
        
        // 处理arp请求
        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)
        default:
            SDLLogger.log("[SDLContext] get invalid packet", for: .debug)
            return .init(inboundBytes: inboundBytes, action: .none)
        }
    }
    
    private func makeARPPlan(layerData: Data, inboundBytes: Int) -> ProcessingPlan {
        // 判断如果收到的是arp请求
        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))")
            }
        } else {
            SDLLogger.log("[SDLContext] get invalid arp packet")
        }
        
        return .init(inboundBytes: inboundBytes, action: .none)
    }
    
    private func makeIPv4Plan(layerData: Data, identityID: UInt32, inboundBytes: Int) -> ProcessingPlan {
        // 有数据是通过出口网关转发的，所有只判断是合法的ip包
        guard let ipPacket = IPPacket(layerData) else {
            return .init(inboundBytes: inboundBytes, action: .none)
        }
        
        switch self.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: IPPacket) -> PolicyPacketContext {
        let ports: (UInt16?, UInt16?)
        switch ipPacket.transportPacket {
        case .tcp(let tcpPacket):
            ports = (tcpPacket.header.srcPort, tcpPacket.header.dstPort)
        case .udp(let udpPacket):
            ports = (udpPacket.srcPort, udpPacket.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
        )
    }
}