package tech.easyflow.ai.easyagentsflow.repository;

import com.easyagents.flow.core.chain.ChainDefinition;
import org.springframework.stereotype.Component;
import org.springframework.transaction.event.TransactionPhase;
import org.springframework.transaction.event.TransactionalEventListener;
import tech.easyflow.ai.easyagentsflow.config.WorkflowRuntimeProperties;
import tech.easyflow.ai.easyagentsflow.event.WorkflowDefinitionChangedEvent;
import tech.easyflow.ai.easyagentsflow.support.PublishedWorkflowDefinitionIds;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.ObjectOutputStream;
import java.time.Duration;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.function.Supplier;

/**
 * 带跨实例版本校验和本地有界 LRU 的工作流编译定义缓存。
 */
@Component
public class WorkflowDefinitionCache {

    private static final int MAX_COMPILE_RETRIES = 3;

    private final WorkflowDefinitionVersionStore versionStore;
    private final int maxEntries;
    private final long maxBytes;
    private final long expireAfterAccessNanos;
    private final Map<String, CacheEntry> entries = new LinkedHashMap<>(16, 0.75F, true);
    private final ConcurrentMap<String, Object> compileLocks = new ConcurrentHashMap<>();
    private long currentBytes;

    /**
     * 创建工作流定义缓存。
     *
     * @param versionStore 跨实例版本令牌仓储
     * @param properties 工作流运行时配置
     */
    public WorkflowDefinitionCache(
            WorkflowDefinitionVersionStore versionStore, WorkflowRuntimeProperties properties) {
        this.versionStore = versionStore;
        this.maxEntries = Math.max(1, properties.getDefinitionCacheMaxEntries());
        this.maxBytes = Math.max(
                1L, properties.getDefinitionCacheMaxBytes());
        Duration expireAfterAccess = properties.getDefinitionCacheExpireAfterAccess();
        this.expireAfterAccessNanos = expireAfterAccess == null
                ? Duration.ofMinutes(30).toNanos()
                : Math.max(1L, expireAfterAccess.toNanos());
    }

    /**
     * 获取已编译定义；缓存未命中时只允许一个线程执行加载与编译。
     *
     * @param definitionId 定义 ID
     * @param loader 定义加载与编译函数
     * @return 已编译工作流定义
     */
    public ChainDefinition get(String definitionId, Supplier<ChainDefinition> loader) {
        String token = versionStore.currentToken(definitionId);
        ChainDefinition cached = getCached(definitionId, token);
        if (cached != null) {
            return cached;
        }

        Object compileLock = compileLocks.computeIfAbsent(definitionId, ignored -> new Object());
        try {
            synchronized (compileLock) {
                for (int attempt = 0; attempt < MAX_COMPILE_RETRIES; attempt++) {
                    token = versionStore.currentToken(definitionId);
                    cached = getCached(definitionId, token);
                    if (cached != null) {
                        return cached;
                    }
                    ChainDefinition compiled = loader.get();
                    String tokenAfterCompile = versionStore.currentToken(definitionId);
                    if (token.equals(tokenAfterCompile)) {
                        put(definitionId, token, compiled);
                        return compiled;
                    }
                }
            }
        } finally {
            compileLocks.remove(definitionId, compileLock);
        }
        throw new IllegalStateException(
                "Workflow definition changed repeatedly while compiling: " + definitionId);
    }

    /**
     * 处理定义变更并更新跨实例版本令牌。
     *
     * @param event 工作流定义变更事件
     */
    @TransactionalEventListener(
            phase = TransactionPhase.AFTER_COMMIT,
            fallbackExecution = true)
    public void onDefinitionChanged(WorkflowDefinitionChangedEvent event) {
        if (event == null || event.workflowId() == null) {
            return;
        }
        versionStore.invalidateWorkflow(event.workflowId());
        synchronized (entries) {
            removeEntry(event.workflowId());
            removeEntry(PublishedWorkflowDefinitionIds.published(
                    event.workflowId()));
        }
    }

    /**
     * 获取仍有效的本地缓存项。
     *
     * @param definitionId 定义 ID
     * @param token 当前版本令牌
     * @return 命中的定义；未命中时返回 null
     */
    private ChainDefinition getCached(String definitionId, String token) {
        long now = System.nanoTime();
        synchronized (entries) {
            CacheEntry entry = entries.get(definitionId);
            if (entry == null) {
                return null;
            }
            if (!entry.token.equals(token) || now - entry.lastAccessNanos > expireAfterAccessNanos) {
                removeEntry(definitionId);
                return null;
            }
            entry.lastAccessNanos = now;
            return entry.definition;
        }
    }

    /**
     * 保存本地缓存并按 LRU 淘汰。
     *
     * @param definitionId 定义 ID
     * @param token 版本令牌
     * @param definition 已编译定义
     */
    private void put(String definitionId, String token, ChainDefinition definition) {
        synchronized (entries) {
            removeEntry(definitionId);
            long weight = serializedSize(definition);
            entries.put(definitionId, new CacheEntry(
                    token,
                    definition,
                    System.nanoTime(),
                    weight));
            currentBytes += weight;
            while (entries.size() > maxEntries
                    || (currentBytes > maxBytes
                    && entries.size() > 1)) {
                String eldestKey = entries.keySet().iterator().next();
                removeEntry(eldestKey);
            }
        }
    }

    /**
     * 删除缓存项并同步维护重量。
     *
     * @param definitionId 定义 ID
     */
    private void removeEntry(String definitionId) {
        CacheEntry removed = entries.remove(definitionId);
        if (removed != null) {
            currentBytes = Math.max(
                    0L, currentBytes - removed.weightBytes);
        }
    }

    /**
     * 使用实际 Java 序列化大小作为缓存重量。
     *
     * @param definition 编译定义
     * @return 序列化字节数
     */
    private long serializedSize(ChainDefinition definition) {
        try (ByteArrayOutputStream output =
                     new ByteArrayOutputStream();
             ObjectOutputStream objects =
                     new ObjectOutputStream(output)) {
            objects.writeObject(definition);
            objects.flush();
            return Math.max(1L, output.size());
        } catch (IOException error) {
            throw new IllegalStateException(
                    "Failed to estimate workflow definition size",
                    error);
        }
    }

    /**
     * 本地定义缓存项。
     */
    private static final class CacheEntry {

        private final String token;
        private final ChainDefinition definition;
        private final long weightBytes;
        private long lastAccessNanos;

        /**
         * 创建本地缓存项。
         *
         * @param token 版本令牌
         * @param definition 已编译定义
         * @param lastAccessNanos 最近访问时间
         */
        private CacheEntry(
                String token,
                ChainDefinition definition,
                long lastAccessNanos,
                long weightBytes) {
            this.token = token;
            this.definition = definition;
            this.lastAccessNanos = lastAccessNanos;
            this.weightBytes = weightBytes;
        }
    }
}
