package tech.easyflow.ai.easyagentsflow.repository;

import com.easyagents.flow.core.chain.ChainState;
import com.easyagents.flow.core.chain.NodeState;
import com.easyagents.flow.core.chain.repository.NodeStateField;
import com.easyagents.flow.core.chain.repository.NodeStateRepository;
import org.springframework.stereotype.Component;
import tech.easyflow.common.cache.VersionedObjectStore;
import tech.easyflow.common.cache.VersionedFields;
import tech.easyflow.common.constant.CacheKey;

import javax.annotation.Resource;
import java.time.Duration;
import java.util.EnumSet;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * 基于 Redis 原子版本存储的节点状态仓储。
 */
@Component
public class NodeStateRepositoryImpl extends BaseRepository implements NodeStateRepository {

    private static final Duration STATE_TTL = Duration.ofDays(3);

    @Resource
    private VersionedObjectStore versionedObjectStore;
    private final ConcurrentMap<String, Boolean> legacyFormats = new ConcurrentHashMap<>();

    /**
     * {@inheritDoc}
     */
    @Override
    public NodeState load(String instanceId, String nodeId) {
        String stateKey = stateKey(instanceId, nodeId);
        VersionedFields snapshot = versionedObjectStore.loadFields(stateKey);
        if (WorkflowStateFields.isFieldFormat(snapshot)) {
            return WorkflowStateFields.decodeNode(snapshot);
        }
        if (snapshot != null) {
            NodeState payloadState = versionedObjectStore.load(stateKey, NodeState.class);
            if (payloadState != null) {
                WorkflowStateFields.normalizeNode(
                        payloadState);
                versionedObjectStore.rewriteAsFields(
                        stateKey,
                        payloadState.getVersion(),
                        WorkflowStateFields.allNodeFields(payloadState),
                        STATE_TTL);
                return payloadState;
            }
        }

        String legacyKey = legacyStateKey(instanceId, nodeId);
        NodeState legacyState = getCache(legacyKey, NodeState.class);
        if (legacyState == null) {
            return null;
        }
        WorkflowStateFields.normalizeNode(
                legacyState);
        if (isLegacyInstance(instanceId)) {
            return legacyState;
        }
        versionedObjectStore.createFieldsIfAbsent(
                stateKey,
                WorkflowStateFields.allNodeFields(legacyState),
                legacyState.getVersion(),
                STATE_TTL);
        VersionedFields migrated = versionedObjectStore.loadFields(stateKey);
        if (!WorkflowStateFields.isFieldFormat(migrated)) {
            throw new IllegalStateException(
                    "Workflow node state migration failed: " + instanceId + "/" + nodeId);
        }
        removeCache(legacyKey);
        return WorkflowStateFields.decodeNode(migrated);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public NodeState create(String instanceId, String nodeId, long chainStateVersion) {
        return create(instanceId, nodeId, chainStateVersion, 0L);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public NodeState create(
            String instanceId,
            String nodeId,
            long chainStateVersion,
            long fencingToken) {
        return create(instanceId, nodeId, chainStateVersion, fencingToken, null, 0L);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public NodeState create(
            String instanceId,
            String nodeId,
            long chainStateVersion,
            long lockFencingToken,
            String claimId,
            long claimGeneration) {
        NodeState existing = load(instanceId, nodeId);
        if (existing != null) {
            return existing;
        }
        NodeState created = new NodeState();
        created.setChainInstanceId(instanceId);
        created.setNodeId(nodeId);
        if (isLegacyInstance(instanceId)) {
            putCache(legacyStateKey(instanceId, nodeId), created);
            return created;
        }
        requireClaimId(claimId, claimGeneration);
        String claimGuardKey = claimGeneration > 0L
                ? executionGuardKey(instanceId, claimId)
                : lockFenceKey(instanceId);
        long effectiveClaimGeneration = claimGeneration > 0L
                ? claimGeneration
                : lockFencingToken;
        boolean createdNow = lockFencingToken > 0L
                ? versionedObjectStore.createFieldsIfAbsent(
                        stateKey(instanceId, nodeId),
                        WorkflowStateFields.allNodeFields(created),
                        created.getVersion(),
                        chainStateKey(instanceId),
                        chainStateVersion,
                        lockFenceKey(instanceId),
                        lockFencingToken,
                        claimGuardKey,
                        effectiveClaimGeneration,
                        STATE_TTL)
                : versionedObjectStore.createFieldsIfAbsent(
                        stateKey(instanceId, nodeId),
                        WorkflowStateFields.allNodeFields(created),
                        created.getVersion(),
                        chainStateKey(instanceId),
                        chainStateVersion,
                        STATE_TTL);
        if (createdNow) {
            return created;
        }
        VersionedFields concurrent = versionedObjectStore.loadFields(
                stateKey(instanceId, nodeId));
        return WorkflowStateFields.isFieldFormat(concurrent)
                ? WorkflowStateFields.decodeNode(concurrent)
                : null;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public boolean tryUpdate(NodeState newState, EnumSet<NodeStateField> fields, long chainStateVersion) {
        return tryUpdate(newState, fields, chainStateVersion, 0L);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public boolean tryUpdate(
            NodeState newState,
            EnumSet<NodeStateField> fields,
            long chainStateVersion,
            long fencingToken) {
        return tryUpdate(
                newState, fields, chainStateVersion, fencingToken, null, 0L);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public boolean tryUpdate(
            NodeState newState,
            EnumSet<NodeStateField> fields,
            long chainStateVersion,
            long lockFencingToken,
            String claimId,
            long claimGeneration) {
        String instanceId = newState.getChainInstanceId();
        if (isLegacyInstance(instanceId)) {
            putCache(legacyStateKey(instanceId, newState.getNodeId()), newState);
            return true;
        }
        long newVersion = newState.getVersion();
        if (newVersion <= 0L) {
            throw new IllegalArgumentException("newState.version must be positive");
        }
        requireClaimId(claimId, claimGeneration);
        String claimGuardKey = claimGeneration > 0L
                ? executionGuardKey(instanceId, claimId)
                : lockFenceKey(instanceId);
        long effectiveClaimGeneration = claimGeneration > 0L
                ? claimGeneration
                : lockFencingToken;
        return lockFencingToken > 0L
                ? versionedObjectStore.compareAndSetFields(
                        stateKey(instanceId, newState.getNodeId()),
                        newVersion - 1L,
                        WorkflowStateFields.nodeFields(newState, fields),
                        newVersion,
                        chainStateKey(instanceId),
                        chainStateVersion,
                        lockFenceKey(instanceId),
                        lockFencingToken,
                        claimGuardKey,
                        effectiveClaimGeneration,
                        STATE_TTL)
                : versionedObjectStore.compareAndSetFields(
                        stateKey(instanceId, newState.getNodeId()),
                        newVersion - 1L,
                        WorkflowStateFields.nodeFields(newState, fields),
                        newVersion,
                        chainStateKey(instanceId),
                        chainStateVersion,
                        STATE_TTL);
    }

    /**
     * 判断实例是否仍由旧 JetCache 状态推进。
     *
     * @param instanceId 实例 ID
     * @return 活跃旧格式实例时为 {@code true}
     */
    private boolean isLegacyInstance(String instanceId) {
        return legacyFormats.computeIfAbsent(instanceId, this::detectLegacyInstance);
    }

    /**
     * 从持久化状态检测实例格式。
     *
     * @param instanceId 实例 ID
     * @return 活跃旧格式实例时为 {@code true}
     */
    private boolean detectLegacyInstance(String instanceId) {
        VersionedFields chainSnapshot = versionedObjectStore.loadFields(
                chainStateKey(instanceId));
        if (WorkflowStateFields.isFieldFormat(chainSnapshot)) {
            return false;
        }
        VersionedFields marker = versionedObjectStore.loadFields(
                CacheKey.CHAIN_STATE_CAS_KEY + "{" + instanceId + "}:format");
        if (marker != null) {
            return false;
        }
        return getCache(CacheKey.CHAIN_CACHE_KEY + instanceId, ChainState.class) != null;
    }

    /**
     * 构建旧节点状态键。
     *
     * @param instanceId 实例 ID
     * @param nodeId 节点 ID
     * @return JetCache 键
     */
    private String legacyStateKey(String instanceId, String nodeId) {
        return CacheKey.NODE_CACHE_KEY + instanceId + ":" + nodeId;
    }

    /**
     * 构建节点状态 CAS key。
     *
     * @param instanceId 工作流实例 ID
     * @param nodeId 节点 ID
     * @return Redis 状态 key
     */
    private String stateKey(String instanceId, String nodeId) {
        return CacheKey.CHAIN_STATE_CAS_KEY + "{" + instanceId + "}:node:" + nodeId;
    }

    /**
     * 构建节点状态提交所依赖的工作流状态 key。
     *
     * @param instanceId 工作流实例 ID
     * @return Redis 工作流状态 key
     */
    private String chainStateKey(String instanceId) {
        return CacheKey.CHAIN_STATE_CAS_KEY + "{" + instanceId + "}:chain";
    }

    /**
     * 构建实例锁 fencing token 键。
     *
     * @param instanceId 工作流实例 ID
     * @return fencing token 键
     */
    private String lockFenceKey(String instanceId) {
        return CacheKey.CHAIN_STATE_CAS_KEY + "{" + instanceId + "}:fence";
    }

    /**
     * 构建实例 fencing token 键。
     *
     * @param instanceId 工作流实例 ID
     * @return fencing token 键
     */
    private String executionGuardKey(String instanceId, String claimId) {
        return CacheKey.CHAIN_STATE_CAS_KEY + "{" + instanceId + "}:claim:" + claimId;
    }

    /**
     * 校验分布式提交所需的 claim ID。
     *
     * @param claimId 触发器 ID
     * @param fencingToken 当前认领 token
     */
    private void requireClaimId(String claimId, long fencingToken) {
        if (fencingToken > 0L && (claimId == null || claimId.trim().isEmpty())) {
            throw new IllegalArgumentException("claimId is required with fencingToken");
        }
    }
}
