//
//  FiveTuple.swift
//  punchnet
//  tcp/udp Flow流管理
//  Created by 安礼成 on 2026/3/10.
//
import Foundation

// MARK: - 五元组 key
struct FlowSession: Hashable {
    let srcIP: UInt32
    let dstIP: UInt32
    let srcPort: UInt16
    let dstPort: UInt16
    let proto: UInt8

    func hash(into hasher: inout Hasher) {
        // 高效组合 hash
        hasher.combine(srcIP)
        hasher.combine(dstIP)
        hasher.combine(UInt32(srcPort) << 16 | UInt32(dstPort))
        hasher.combine(proto)
    }

    static func ==(lhs: Self, rhs: Self) -> Bool {
        return lhs.srcIP == rhs.srcIP &&
               lhs.dstIP == rhs.dstIP &&
               lhs.srcPort == rhs.srcPort &&
               lhs.dstPort == rhs.dstPort &&
               lhs.proto == rhs.proto
    }
    
    func reverse() -> FlowSession {
        return FlowSession(
            srcIP: dstIP,
            dstIP: srcIP,
            srcPort: dstPort,
            dstPort: srcPort,
            proto: proto
        )
    }
    
}

// MARK: - 会话管理器
final class SDLFlowSessionManager {
    private var sessions: [FlowSession: TimeInterval] = [:]
    private let lock = NSLock()
    private let sessionTimeout: TimeInterval

    /// - Parameter sessionTimeout: 会话闲置多久（秒）被清理
    init(sessionTimeout: TimeInterval = 300) {
        self.sessionTimeout = sessionTimeout
    }

    // 插入或更新会话
    func updateSession(_ key: FlowSession) {
        lock.lock()
        defer {
            lock.unlock()
        }
        sessions[key] = Date().timeIntervalSince1970 + sessionTimeout
    }
    
    // 查找会话
    func hasSession(_ key: FlowSession) -> Bool {
        lock.lock()
        defer {
            lock.unlock()
        }
        
        if let expireTs = sessions[key] {
            if expireTs >= Date().timeIntervalSince1970 {
                return true
            }
            self.sessions.removeValue(forKey: key)
        }
        
        return false
    }

    // 删除会话
    func removeSession(_ key: FlowSession) {
        lock.lock()
        defer {
            lock.unlock()
        }
        
        sessions.removeValue(forKey: key)
    }

    // 清理过期会话
    func cleanupExpiredSessions() {
        lock.lock()
        defer {
            lock.unlock()
        }
        
        let now = Date().timeIntervalSince1970
        self.sessions = self.sessions.filter { $0.value >= now }
    }

    // 返回当前会话数（调试/统计用）
    var count: Int {
        lock.lock()
        defer {
            lock.unlock()
        }
        return sessions.count
    }
    
}

extension IPPacket {
    
    func flowSession() -> FlowSession? {
        switch self.transportPacket {
        case .tcp(let tcpPacket):
            return FlowSession(srcIP: header.source, dstIP: header.destination, srcPort: tcpPacket.header.srcPort, dstPort: tcpPacket.header.dstPort, proto: header.proto)
        case .udp(let udpPacket):
            return FlowSession(srcIP: header.source, dstIP: header.destination, srcPort: udpPacket.srcPort, dstPort: udpPacket.dstPort, proto: header.proto)
        default:
            return nil
        }
    }
}