package tech.easyflow.ai.easyagentsflow.repository;

import com.easyagents.flow.core.chain.ChainDefinition;
import com.easyagents.flow.core.chain.repository.ChainDefinitionSnapshotRepository;
import org.springframework.stereotype.Component;
import tech.easyflow.common.cache.VersionedObjectStore;
import tech.easyflow.common.constant.CacheKey;

import javax.annotation.Resource;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.ObjectOutputStream;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.time.Duration;
import java.util.Collections;
import java.util.HexFormat;
import java.util.List;
import java.util.Map;
import java.util.WeakHashMap;

/**
 * 基于 Redis 版本对象存储的工作流实例定义快照仓储。
 */
@Component
public class ChainDefinitionSnapshotRepositoryImpl extends BaseRepository
        implements ChainDefinitionSnapshotRepository {

    /**
     * 快照由工作流终态显式删除；长 TTL 仅用于异常中断后的孤儿兜底清理。
     */
    private static final Duration SNAPSHOT_TTL = Duration.ofDays(7);
    private static final Map<ChainDefinition, String> CONTENT_HASH_CACHE =
            Collections.synchronizedMap(new WeakHashMap<>());

    @Resource
    private VersionedObjectStore versionedObjectStore;

    /**
     * {@inheritDoc}
     */
    @Override
    public void save(String instanceId, ChainDefinition definition) {
        String contentHash = contentHash(definition);
        String immutableContentKey = contentKey(contentHash);
        versionedObjectStore.createIfAbsent(
                immutableContentKey,
                definition,
                0L,
                SNAPSHOT_TTL);
        // 每个新实例都延长不可变内容寿命，保证引用 TTL 内不会悬空。
        versionedObjectStore.refreshExpirations(
                List.of(immutableContentKey), SNAPSHOT_TTL);
        versionedObjectStore.createIfAbsent(
                stateKey(instanceId),
                contentHash,
                0L,
                SNAPSHOT_TTL);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public ChainDefinition load(String instanceId) {
        Object reference = versionedObjectStore.load(
                stateKey(instanceId), Object.class);
        if (reference instanceof String) {
            ChainDefinition snapshot = versionedObjectStore.load(
                    contentKey((String) reference),
                    ChainDefinition.class);
            if (snapshot != null) {
                return snapshot;
            }
        } else if (reference instanceof ChainDefinition) {
            // 兼容 XL12 之前按实例保存完整定义的运行中实例。
            return (ChainDefinition) reference;
        }
        ChainDefinition legacy = getCache(legacyKey(instanceId), ChainDefinition.class);
        if (legacy != null) {
            save(instanceId, legacy);
        }
        return legacy;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void remove(String instanceId) {
        versionedObjectStore.deleteAll(List.of(stateKey(instanceId)));
        removeCache(legacyKey(instanceId));
    }

    /**
     * 构建定义快照缓存键。
     *
     * @param instanceId 工作流实例 ID
     * @return 缓存键
     */
    private String legacyKey(String instanceId) {
        return CacheKey.CHAIN_DEFINITION_SNAPSHOT_CACHE_KEY + instanceId;
    }

    /**
     * 构建不会随短期运行状态 TTL 漂移的快照键。
     *
     * @param instanceId 工作流实例 ID
     * @return Redis 快照键
     */
    private String stateKey(String instanceId) {
        return CacheKey.CHAIN_STATE_CAS_KEY + "{" + instanceId + "}:definition";
    }

    /**
     * 构建按定义内容复用的快照键。
     *
     * @param contentHash 定义内容摘要
     * @return Redis 内容键
     */
    private String contentKey(String contentHash) {
        return CacheKey.CHAIN_DEFINITION_SNAPSHOT_CACHE_KEY
                + "content:"
                + contentHash;
    }

    /**
     * 计算定义序列化内容摘要；同一编译定义对象仅计算一次。
     *
     * @param definition 编译后的定义
     * @return SHA-256 十六进制摘要
     */
    private String contentHash(ChainDefinition definition) {
        if (definition == null) {
            throw new IllegalArgumentException(
                    "Chain definition required");
        }
        String cached = CONTENT_HASH_CACHE.get(definition);
        if (cached != null) {
            return cached;
        }
        try (ByteArrayOutputStream output =
                     new ByteArrayOutputStream();
             ObjectOutputStream objects =
                     new ObjectOutputStream(output)) {
            objects.writeObject(definition);
            objects.flush();
            String calculated = HexFormat.of().formatHex(
                    MessageDigest.getInstance("SHA-256")
                            .digest(output.toByteArray()));
            CONTENT_HASH_CACHE.put(definition, calculated);
            return calculated;
        } catch (IOException | NoSuchAlgorithmException error) {
            throw new IllegalStateException(
                    "Failed to hash chain definition", error);
        }
    }
}
