package tech.easyflow.ai.easyagentsflow.repository;

import com.easyagents.flow.core.chain.ChainState;
import com.easyagents.flow.core.chain.repository.ChainLock;
import com.easyagents.flow.core.chain.repository.ChainStateField;
import com.easyagents.flow.core.chain.repository.ChainStateRepository;
import org.springframework.stereotype.Component;
import tech.easyflow.common.cache.RedisLockExecutor;
import tech.easyflow.common.cache.VersionedFields;
import tech.easyflow.common.cache.VersionedObjectStore;
import tech.easyflow.common.constant.CacheKey;

import javax.annotation.Resource;
import java.io.Serializable;
import java.time.Duration;
import java.util.EnumSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 基于 Redis 字段化 CAS、认领守卫和短期实例锁的工作流状态仓储。
 */
@Component
public class ChainStateRepositoryImpl extends BaseRepository implements ChainStateRepository {

    private static final Duration STATE_TTL = Duration.ofDays(3);
    private static final Duration MIGRATION_MARKER_TTL = Duration.ofDays(4);
    private static final Duration FENCING_COUNTER_TTL = Duration.ofDays(4);
    private static final Duration MIN_LOCK_LEASE = Duration.ofSeconds(30);
    private static final ScheduledThreadPoolExecutor LOCK_RENEW_EXECUTOR =
            createLockRenewExecutor();
    private final Set<String> legacyInstances = java.util.concurrent.ConcurrentHashMap.newKeySet();

    @Resource
    private RedisLockExecutor redisLockExecutor;
    @Resource
    private VersionedObjectStore versionedObjectStore;

    /**
     * 创建会主动移除已取消任务的锁续期线程池。
     *
     * <p>绝大多数实例锁仅持有数毫秒；开启 remove-on-cancel 可避免高吞吐场景下，
     * 已取消的十秒延迟续期任务在队列中短时堆积。</p>
     *
     * @return 小型多线程锁续期调度器
     */
    private static ScheduledThreadPoolExecutor createLockRenewExecutor() {
        ScheduledThreadPoolExecutor executor =
                new ScheduledThreadPoolExecutor(2, new LockRenewThreadFactory());
        executor.setRemoveOnCancelPolicy(true);
        executor.setExecuteExistingDelayedTasksAfterShutdownPolicy(false);
        return executor;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public ChainState load(String instanceId) {
        String stateKey = stateKey(instanceId);
        VersionedFields snapshot = versionedObjectStore.loadFields(stateKey);
        if (WorkflowStateFields.isFieldFormat(snapshot)) {
            legacyInstances.remove(instanceId);
            return WorkflowStateFields.decodeChain(snapshot);
        }
        if (snapshot != null) {
            ChainState payloadState = versionedObjectStore.load(stateKey, ChainState.class);
            if (payloadState != null) {
                rewritePayloadState(stateKey, payloadState);
                return payloadState;
            }
        }
        if (hasMigrationMarker(instanceId)) {
            legacyInstances.remove(instanceId);
            return null;
        }

        String legacyKey = CacheKey.CHAIN_CACHE_KEY + instanceId;
        ChainState legacyState = getCache(legacyKey, ChainState.class);
        if (legacyState == null) {
            return null;
        }
        // 活跃旧实例继续沿用旧写路径，避免滚动升级期间两个格式同时推进。
        if (legacyState.getStatus() == null || !legacyState.getStatus().isTerminal()) {
            legacyInstances.add(instanceId);
            return legacyState;
        }
        migrateTerminalLegacyState(instanceId, legacyKey, legacyState);
        VersionedFields migrated = versionedObjectStore.loadFields(stateKey);
        if (!WorkflowStateFields.isFieldFormat(migrated)) {
            throw new IllegalStateException("Workflow state migration failed: " + instanceId);
        }
        return WorkflowStateFields.decodeChain(migrated);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public Long loadVersion(String instanceId) {
        if (!legacyInstances.contains(instanceId)) {
            Long version = versionedObjectStore.loadVersion(stateKey(instanceId));
            if (version != null) {
                return version;
            }
        }
        ChainState state = load(instanceId);
        return state == null ? null : state.getVersion();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public ChainState create(String instanceId) {
        ChainState existing = load(instanceId);
        if (existing != null) {
            return existing;
        }
        ChainState created = new ChainState();
        created.setInstanceId(instanceId);
        if (versionedObjectStore.createFieldsIfAbsent(
                stateKey(instanceId),
                WorkflowStateFields.allChainFields(created),
                created.getVersion(),
                STATE_TTL)) {
            touchMigrationMarker(instanceId);
            return created;
        }
        VersionedFields concurrent = versionedObjectStore.loadFields(stateKey(instanceId));
        if (!WorkflowStateFields.isFieldFormat(concurrent)) {
            throw new IllegalStateException("Unable to initialize workflow state: " + instanceId);
        }
        return WorkflowStateFields.decodeChain(concurrent);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public boolean tryUpdate(ChainState newState, EnumSet<ChainStateField> fields) {
        return tryUpdate(newState, fields, 0L);
    }

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

    /**
     * {@inheritDoc}
     */
    @Override
    public boolean tryUpdate(
            ChainState newState,
            EnumSet<ChainStateField> fields,
            long lockFencingToken,
            String claimId,
            long claimGeneration) {
        String instanceId = newState.getInstanceId();
        if (legacyInstances.contains(instanceId)) {
            String legacyKey = CacheKey.CHAIN_CACHE_KEY + instanceId;
            putCache(legacyKey, newState);
            if (newState.getStatus() != null && newState.getStatus().isTerminal()) {
                migrateTerminalLegacyState(instanceId, legacyKey, newState);
            }
            return true;
        }
        long newVersion = newState.getVersion();
        if (newVersion <= 0L) {
            throw new IllegalArgumentException("newState.version must be positive");
        }
        requireClaimId(claimId, claimGeneration);
        boolean updated;
        if (lockFencingToken > 0L) {
            String claimGuardKey = claimGeneration > 0L
                    ? executionGuardKey(instanceId, claimId)
                    : lockFenceKey(instanceId);
            long effectiveClaimGeneration = claimGeneration > 0L
                    ? claimGeneration
                    : lockFencingToken;
            updated = versionedObjectStore.compareAndSetFieldsAndRefresh(
                        stateKey(instanceId),
                        newVersion - 1L,
                        WorkflowStateFields.chainFields(newState, fields),
                        newVersion,
                        lockFenceKey(instanceId),
                        lockFencingToken,
                        claimGuardKey,
                        effectiveClaimGeneration,
                        STATE_TTL,
                        markerKey(instanceId),
                        MIGRATION_MARKER_TTL);
        } else {
            updated = versionedObjectStore.compareAndSetFields(
                        stateKey(instanceId),
                        newVersion - 1L,
                        WorkflowStateFields.chainFields(newState, fields),
                        newVersion,
                        STATE_TTL);
        }
        return updated;
    }

    /**
     * 将当前 Redis Hash 中的旧完整 payload 原地改写为字段化状态。
     *
     * @param stateKey 状态键
     * @param state 旧完整状态
     */
    private void rewritePayloadState(String stateKey, ChainState state) {
        boolean rewritten = versionedObjectStore.rewriteAsFields(
                stateKey,
                state.getVersion(),
                WorkflowStateFields.allChainFields(state),
                STATE_TTL);
        if (rewritten) {
            touchMigrationMarker(state.getInstanceId());
        }
    }

    /**
     * 在旧实例终态后完成受控迁移并留下 tombstone。
     *
     * @param instanceId 实例 ID
     * @param legacyKey 旧 JetCache 键
     * @param state 终态状态
     */
    private void migrateTerminalLegacyState(
            String instanceId, String legacyKey, ChainState state) {
        boolean created = versionedObjectStore.createFieldsIfAbsent(
                stateKey(instanceId),
                WorkflowStateFields.allChainFields(state),
                state.getVersion(),
                STATE_TTL);
        VersionedFields existing = versionedObjectStore.loadFields(stateKey(instanceId));
        if (!created && !WorkflowStateFields.isFieldFormat(existing)) {
            throw new IllegalStateException("Workflow state migration conflict: " + instanceId);
        }
        touchMigrationMarker(instanceId);
        removeCache(legacyKey);
        legacyInstances.remove(instanceId);
    }

    /**
     * 判断实例是否已经切换到字段化格式。
     *
     * @param instanceId 实例 ID
     * @return 已存在迁移 tombstone 时为 {@code true}
     */
    private boolean hasMigrationMarker(String instanceId) {
        return versionedObjectStore.loadFields(markerKey(instanceId)) != null;
    }

    /**
     * 创建或刷新迁移 tombstone，防止新状态过期后复活旧缓存。
     *
     * @param instanceId 实例 ID
     */
    private void touchMigrationMarker(String instanceId) {
        Map<String, Serializable> marker = Map.of(
                WorkflowStateFields.FORMAT_FIELD, WorkflowStateFields.FORMAT_VERSION);
        if (!versionedObjectStore.createFieldsIfAbsent(
                markerKey(instanceId), marker, 0L, MIGRATION_MARKER_TTL)) {
            versionedObjectStore.compareAndSetFields(
                    markerKey(instanceId), 0L, marker, 0L, MIGRATION_MARKER_TTL);
        }
    }

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

    /**
     * 构建状态格式 tombstone 键。
     *
     * @param instanceId 实例 ID
     * @return Redis marker 键
     */
    private String markerKey(String instanceId) {
        return CacheKey.CHAIN_STATE_CAS_KEY + "{" + instanceId + "}:format";
    }

    /**
     * 构建本次触发器认领的执行守卫键。
     *
     * @param instanceId 工作流实例 ID
     * @param claimId 触发器 ID
     * @return 执行守卫键
     */
    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");
        }
    }

    /**
     * 获取工作流实例级 Redis 分布式锁。
     *
     * @param instanceId 工作流实例 ID
     * @param timeout 等待锁的最大时间
     * @param unit 时间单位
     * @return 分布式锁句柄
     */
    @Override
    public ChainLock getLock(String instanceId, long timeout, TimeUnit unit) {
        if (instanceId == null || instanceId.trim().isEmpty()) {
            throw new IllegalArgumentException("instanceId must not be blank");
        }
        if (redisLockExecutor == null) {
            throw new IllegalStateException("RedisLockExecutor is not configured");
        }
        Duration waitTimeout = Duration.ofMillis(Math.max(1L, unit.toMillis(timeout)));
        Duration leaseTimeout = waitTimeout.compareTo(MIN_LOCK_LEASE) > 0 ? waitTimeout : MIN_LOCK_LEASE;
        RedisLockExecutor.LockHandle handle = redisLockExecutor.tryAcquireFenced(
                CacheKey.CHAIN_LOCK_KEY + "{" + instanceId + "}",
                lockFenceKey(instanceId),
                waitTimeout,
                leaseTimeout,
                FENCING_COUNTER_TTL);
        if (handle == null) {
            return new RedisChainLock(null, 0L);
        }
        return new RedisChainLock(handle, handle.getFencingToken());
    }

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

    /**
     * Redis 工作流锁适配器。
     */
    private static final class RedisChainLock implements ChainLock {

        private final RedisLockExecutor.LockHandle handle;
        private final long fencingToken;
        private final AtomicBoolean valid = new AtomicBoolean(true);
        private final ScheduledFuture<?> renewTask;

        /**
         * 创建锁适配器。
         *
         * @param handle Redis 锁句柄；为空表示未获取
         * @param fencingToken 本次锁持有期的 fencing token
         */
        private RedisChainLock(
                RedisLockExecutor.LockHandle handle, long fencingToken) {
            this.handle = handle;
            this.fencingToken = fencingToken;
            this.renewTask = handle == null ? null : LOCK_RENEW_EXECUTOR.scheduleWithFixedDelay(
                    () -> {
                        if (!handle.renew()) {
                            valid.set(false);
                        }
                    },
                    MIN_LOCK_LEASE.toMillis() / 3L,
                    MIN_LOCK_LEASE.toMillis() / 3L,
                    TimeUnit.MILLISECONDS);
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public boolean isAcquired() {
            return handle != null && valid.get();
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public boolean isValid() {
            return isAcquired();
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public long getFencingToken() {
            return fencingToken;
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public void close() {
            valid.set(false);
            if (renewTask != null) {
                renewTask.cancel(false);
            }
            if (handle != null) {
                handle.release();
            }
        }
    }

    /**
     * 创建守护型工作流状态锁续期线程。
     */
    private static final class LockRenewThreadFactory implements ThreadFactory {

        /**
         * {@inheritDoc}
         */
        @Override
        public Thread newThread(Runnable runnable) {
            Thread thread = new Thread(runnable, "workflow-state-lock-renew");
            thread.setDaemon(true);
            return thread;
        }
    }
}
