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

import Foundation

final class SDLHoleDataProcessor {
    enum ProcessingAction {
        case sendARPReply(dstMac: Data, data: Data)
        case appendARP(ip: UInt32, mac: Data)
        case writeToTun(packetData: Data, identityID: UInt32)
        case requestPolicy(srcIdentityID: UInt32)
        case none
    }
    
    struct ProcessingPlan {
        let inboundBytes: Int
        let action: ProcessingAction
    }
    
    private let networkAddress: SDLConfiguration.NetworkAddress
    private let dataCipher: CCDataCipher?
    private let snapshotPublisher: SnapshotPublisher<IdentitySnapshot>
    private let flowSessionManager: SDLFlowSessionManager
    
    init(networkAddress: SDLConfiguration.NetworkAddress,
         dataCipher: CCDataCipher?,
         snapshotPublisher: SnapshotPublisher<IdentitySnapshot>,
         flowSessionManager: SDLFlowSessionManager) {
        self.networkAddress = networkAddress
        self.dataCipher = dataCipher
        self.snapshotPublisher = snapshotPublisher
        self.flowSessionManager = flowSessionManager
    }
    
    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)
        }
        
        // 检查权限逻辑
        let identitySnapshot = self.snapshotPublisher.current()
        let ruleMap = identitySnapshot.lookup(identityID)
        
        if true || self.checkPolicy(ipPacket: ipPacket, ruleMap: ruleMap) {
            return .init(
                inboundBytes: inboundBytes,
                action: .writeToTun(packetData: ipPacket.data, identityID: identityID)
            )
        }
        
        return .init(
            inboundBytes: inboundBytes,
            action: .requestPolicy(srcIdentityID: identityID)
        )
    }
    
    private func checkPolicy(ipPacket: IPPacket, ruleMap: IdentityRuleMap?) -> Bool {
        // 进来的数据反转一下，然后再处理
        if let reverseFlowSession = ipPacket.flowSession()?.reverse(),
           self.flowSessionManager.hasSession(reverseFlowSession) {
            self.flowSessionManager.updateSession(reverseFlowSession)
            return true
        }
        
        // 检查权限逻辑
        let proto = ipPacket.header.proto
        // 优先判断访问规则
        switch ipPacket.transportPacket {
        case .tcp(let tcpPacket):
            if let ruleMap, ruleMap.isAllow(proto: proto, port: tcpPacket.header.dstPort) {
                return true
            }
        case .udp(let udpPacket):
            if let ruleMap, ruleMap.isAllow(proto: proto, port: udpPacket.dstPort) {
                return true
            }
        case .icmp(_):
            return true
        default:
            return false
        }
        
        return false
    }
}