package com.cmvr.test.flow.runtime.engine.support;

import org.springframework.stereotype.Component;

import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicBoolean;

@Component
public class FlowSchedulerCache {

    // instId_nodeId 或 instId_nodeId_loopIteration
    private final ConcurrentMap<String, AtomicBoolean> scheduled = new ConcurrentHashMap<>();
    private final ConcurrentMap<String, Set<String>> completedPreMap = new ConcurrentHashMap<>();

    public boolean trySchedule(String key) {
        AtomicBoolean isScheduled = scheduled.computeIfAbsent(key, k -> new AtomicBoolean(false));
        return isScheduled.compareAndSet(false, true);
    }

    public Set<String> getOrInitCompletedPre(String key) {
        return completedPreMap.computeIfAbsent(key, k -> ConcurrentHashMap.newKeySet());
    }

    public void clearContextCache(String instId) {
        scheduled.keySet().removeIf(k -> k.startsWith(instId + "_"));
        completedPreMap.keySet().removeIf(k -> k.startsWith(instId + "_"));
    }
}
