//
//  IdentityStore.swift
//  punchnet
//
//  Created by 安礼成 on 2026/2/5.
//
import Foundation
import NIO

final class IdentitySession {
    var version: UInt32
    var totalNum: UInt32
    private var parts: [UInt32: SDLPolicyResponse] = [:]
    
    init(part: SDLPolicyResponse) {
        self.version = part.version
        self.totalNum = part.totalNum
        self.parts[part.index] = part
    }
    
    func merge(part: SDLPolicyResponse) {
        if part.version < version {
            // 低版本数据丢弃
        } else if part.version == version {
            self.parts[part.index] = part
        } else {
            self.parts.removeAll()
            self.parts[part.index] = part
        }
    }
    
    func process() -> Data? {
        // parts是连续的，从0开始，并且数量等于total_num
        let indexs = parts.keys.sorted().map { UInt32($0) }
        guard indexs.count == self.totalNum && isContinuousFromZero(indexs: indexs) else {
            return nil
        }
        
        var rulesData: Data = Data()
        for i in 0..<totalNum {
            if let part = self.parts[i] {
                rulesData.append(part.rules)
            }
        }
        SDLLogger.shared.log("[IdentitySession] get a completed rules: \(rulesData.count)")
        
        return rulesData
    }
    
    private func isContinuousFromZero(indexs: [UInt32]) -> Bool {
        guard !indexs.isEmpty else {
            return false
        }
        
        return indexs.enumerated().allSatisfy { idx, value in
            idx == value
        }
    }
    
}

actor IdentityStore {
    typealias IdentityID = UInt32
    
    nonisolated private let alloctor = ByteBufferAllocator()
    
    private let publisher: SnapshotPublisher<IdentitySnapshot>
    private var identityMap: [IdentityID: IdentityRuleMap] = [:]
    private var sessions: [IdentityID: IdentitySession] = [:]
    
    init(publisher: SnapshotPublisher<IdentitySnapshot>) {
        self.publisher = publisher
    }
    
    func apply(policyResponse: SDLPolicyResponse) {
        let id = policyResponse.srcIdentityID
        let session = self.sessions[id, default: IdentitySession(part: policyResponse)]
        session.merge(part: policyResponse)
        
        // 判断一下是否接受完成
        if let rulesData = session.process() {
            var buffer = alloctor.buffer(bytes: rulesData)
            var ruleMap: [UInt8: [UInt16: Bool]] = [:]
            while true {
                guard let proto = buffer.readInteger(endianness: .big, as: UInt8.self),
                      let port = buffer.readInteger(endianness: .big, as: UInt16.self) else {
                    break
                }
                ruleMap[proto, default: [:]][port] = true
            }
            self.identityMap[id] = IdentityRuleMap(ruleMap: ruleMap)
            
            // 删除当前的session信息
            self.sessions.removeValue(forKey: id)
            
            SDLLogger.shared.log("[IdentitySession] get compile Snapshot rules nums: \(self.identityMap[id]?.ruleMap.count), success: \(self.identityMap[id]?.isAllow(proto: 1, port: 80))")
            
            // 发布新的快照信息
            let snapshot = compileSnapshot()
            publisher.publish(snapshot)
        } else {
            self.sessions[id] = session
        }
    }
    
    private func compileSnapshot() -> IdentitySnapshot {
        return IdentitySnapshot(identityMap: identityMap)
    }

}