package tech.easyflow.ai.easyagentsflow.repository;

import com.easyagents.flow.core.chain.repository.LoopResultReference;
import com.easyagents.flow.core.chain.repository.LoopInputReference;
import com.easyagents.flow.core.chain.runtime.TriggerClaimLostException;
import org.junit.Assert;
import org.junit.Test;
import tech.easyflow.common.cache.VersionedObjectStore;

import java.lang.reflect.Field;
import java.io.Serializable;
import java.time.Duration;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyList;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

/**
 * 循环结果分块仓储测试。
 */
public class LoopResultRepositoryImplTest {

    /**
     * 验证跨多个分块的结果顺序、完整性及幂等重放。
     */
    @Test
    public void shouldPreserveOrderingAcrossChunkBoundaries() {
        InMemoryLoopResultRepository repository = new InMemoryLoopResultRepository();
        String resultId = "loop-result";
        int iterations = LoopResultRepositoryImpl.CHUNK_SIZE * 2 + 1;

        for (int index = 0; index < iterations; index++) {
            repository.append(resultId, index, Map.of(
                    "index", index,
                    "value", "value-" + index));
        }
        repository.append(resultId, iterations - 1, Map.of(
                "index", iterations - 1,
                "value", "value-" + (iterations - 1)));

        Map<String, Object> result =
                repository.load(resultId, iterations, List.of("index", "value"));
        Assert.assertEquals(iterations, ((List<?>) result.get("index")).size());
        Assert.assertEquals(0, ((List<?>) result.get("index")).get(0));
        Assert.assertEquals(iterations - 1, ((List<?>) result.get("index")).get(iterations - 1));
        Assert.assertEquals("value-128", ((List<?>) result.get("value")).get(128));
    }

    /**
     * 验证同一轮次写入不同结果时拒绝覆盖。
     */
    @Test(expected = IllegalStateException.class)
    public void shouldRejectConflictingReplay() {
        InMemoryLoopResultRepository repository = new InMemoryLoopResultRepository();
        repository.append("loop-result", 0, Map.of("value", "first"));
        repository.append("loop-result", 0, Map.of("value", "changed"));
    }

    /**
     * 验证热状态只保存轻量引用，业务读取边界仍还原为原有列表结构。
     */
    @Test
    public void shouldResolveLightweightReferenceAtReadBoundary() {
        InMemoryLoopResultRepository repository = new InMemoryLoopResultRepository();
        String resultId = "instance:loop-result";
        repository.append(resultId, 0, Map.of("value", "first"));
        repository.append(resultId, 1, Map.of("value", "second"));

        Map<String, Object> references = repository.references(
                resultId, 2, List.of("value"));
        Assert.assertTrue(references.get("value") instanceof LoopResultReference);

        @SuppressWarnings("unchecked")
        Map<String, Object> resolved =
                (Map<String, Object>) repository.resolveReferences(references);
        Assert.assertEquals(List.of("first", "second"), resolved.get("value"));
    }

    /**
     * 验证没有声明输出的长循环跨分块时仍会续期输入生命周期。
     *
     * @throws Exception 测试依赖注入失败时抛出
     */
    @Test
    public void shouldRefreshInputChunksWhenLoopHasNoOutputs() throws Exception {
        LoopResultRepositoryImpl repository = new LoopResultRepositoryImpl();
        VersionedObjectStore store = mock(VersionedObjectStore.class);
        Field field = LoopResultRepositoryImpl.class.getDeclaredField(
                "versionedObjectStore");
        field.setAccessible(true);
        field.set(repository, store);
        when(store.load(any(String.class), eq(Integer.class))).thenReturn(256);
        when(store.createIfAbsent(
                anyString(),
                any(Serializable.class),
                anyLong(),
                anyString(),
                anyLong(),
                anyString(),
                anyLong(),
                any(Duration.class))).thenReturn(true);

        repository.append(
                "instance",
                1L,
                "claim",
                1L,
                "instance:loop",
                LoopResultRepositoryImpl.CHUNK_SIZE,
                Map.of());

        verify(store).refreshExpirations(anyList(), any(Duration.class));
    }

    /**
     * 验证同一输入分块内的多轮读取只访问一次底层对象存储。
     *
     * @throws Exception 测试依赖注入失败时抛出
     */
    @Test
    public void shouldLoadEachInputChunkOnlyOnce() throws Exception {
        LoopResultRepositoryImpl repository = new LoopResultRepositoryImpl();
        VersionedObjectStore store = mock(VersionedObjectStore.class);
        Field field = LoopResultRepositoryImpl.class.getDeclaredField(
                "versionedObjectStore");
        field.setAccessible(true);
        field.set(repository, store);
        List<Object> values = java.util.stream.IntStream.range(0, 128)
                .boxed()
                .map(value -> (Object) value)
                .toList();
        when(store.load(anyString(), eq(List.class))).thenReturn(values);

        Assert.assertEquals(0, repository.loadInputItem("instance:loop", 0));
        Assert.assertEquals(127, repository.loadInputItem("instance:loop", 127));

        verify(store, times(1)).load(anyString(), eq(List.class));
    }

    /**
     * 验证调用方修改已读取的可变输入时不会污染活动分块缓存。
     *
     * @throws Exception 测试依赖注入失败时抛出
     */
    @Test
    public void shouldIsolateMutableInputValuesFromActiveCache()
            throws Exception {
        LoopResultRepositoryImpl repository =
                new LoopResultRepositoryImpl();
        VersionedObjectStore store =
                mock(VersionedObjectStore.class);
        Field field = LoopResultRepositoryImpl.class
                .getDeclaredField("versionedObjectStore");
        field.setAccessible(true);
        field.set(repository, store);
        Map<String, Object> persisted =
                new LinkedHashMap<>();
        persisted.put("name", "original");
        when(store.load(anyString(), eq(List.class)))
                .thenReturn(List.of(persisted));

        @SuppressWarnings("unchecked")
        Map<String, Object> first =
                (Map<String, Object>)
                        repository.loadInputItem(
                                "instance:mutable-input",
                                0);
        first.put("name", "changed");
        @SuppressWarnings("unchecked")
        Map<String, Object> second =
                (Map<String, Object>)
                        repository.loadInputItem(
                                "instance:mutable-input",
                                0);

        Assert.assertEquals(
                "original", second.get("name"));
        verify(store, times(1)).load(
                anyString(), eq(List.class));
    }

    /**
     * 验证完整输入还原按分块批量读取，并保持不同调用方的可变值隔离。
     *
     * @throws Exception 测试依赖注入失败时抛出
     */
    @Test
    @SuppressWarnings("unchecked")
    public void shouldBulkLoadMutableInputChunks()
            throws Exception {
        LoopResultRepositoryImpl repository =
                new LoopResultRepositoryImpl();
        VersionedObjectStore store =
                mock(VersionedObjectStore.class);
        Field field = LoopResultRepositoryImpl.class
                .getDeclaredField("versionedObjectStore");
        field.setAccessible(true);
        field.set(repository, store);
        int itemCount =
                LoopResultRepositoryImpl.CHUNK_SIZE
                        * 2 + 1;
        when(store.loadAll(
                anyList(), eq(List.class)))
                .thenAnswer(invocation -> {
                    List<List<Object>> chunks =
                            new java.util.ArrayList<>();
                    for (int chunkIndex = 0;
                         chunkIndex < 3;
                         chunkIndex++) {
                        int chunkSize = chunkIndex < 2
                                ? LoopResultRepositoryImpl.CHUNK_SIZE
                                : 1;
                        List<Object> chunk =
                                new java.util.ArrayList<>();
                        for (int offset = 0;
                             offset < chunkSize;
                             offset++) {
                            Map<String, Object> value =
                                    new LinkedHashMap<>();
                            value.put(
                                    "index",
                                    chunkIndex
                                            * LoopResultRepositoryImpl.CHUNK_SIZE
                                            + offset);
                            chunk.add(value);
                        }
                        chunks.add(chunk);
                    }
                    return chunks;
                });
        LoopInputReference reference =
                new LoopInputReference(
                        "instance:bulk-input",
                        itemCount);

        List<Object> first =
                repository.loadInput(reference);
        ((Map<String, Object>) first.get(0))
                .put("index", -1);
        List<Object> second =
                repository.loadInput(reference);

        Assert.assertEquals(
                itemCount, second.size());
        Assert.assertEquals(
                0,
                ((Map<String, Object>) second.get(0))
                        .get("index"));
        verify(store, times(2)).loadAll(
                anyList(), eq(List.class));
        verify(store, times(0)).load(
                anyString(), eq(List.class));
    }

    /**
     * 验证连续循环输出命中活动分块缓存时不重复读取 Redis 对象。
     *
     * @throws Exception 测试依赖注入失败时抛出
     */
    @Test
    public void shouldReuseActiveOutputChunkAfterSuccessfulCommit()
            throws Exception {
        LoopResultRepositoryImpl repository =
                new LoopResultRepositoryImpl();
        VersionedObjectStore store =
                mock(VersionedObjectStore.class);
        Field field = LoopResultRepositoryImpl.class
                .getDeclaredField("versionedObjectStore");
        field.setAccessible(true);
        field.set(repository, store);
        when(store.load(
                anyString(),
                eq(LoopResultRepositoryImpl
                        .LoopResultChunk.class)))
                .thenReturn(null);
        when(store.createIfAbsent(
                anyString(),
                any(Serializable.class),
                anyLong(),
                anyString(),
                anyLong(),
                anyString(),
                anyLong(),
                any(Duration.class)))
                .thenReturn(true);
        when(store.compareAndSet(
                anyString(),
                anyLong(),
                any(Serializable.class),
                anyLong(),
                anyString(),
                anyLong(),
                anyString(),
                anyLong(),
                any(Duration.class)))
                .thenReturn(true);

        repository.append(
                "instance",
                1L,
                "claim",
                1L,
                "instance:cached-output",
                0,
                Map.of("value", "first"));
        repository.append(
                "instance",
                1L,
                "claim",
                1L,
                "instance:cached-output",
                1,
                Map.of("value", "second"));

        verify(store, times(1)).load(
                anyString(),
                eq(LoopResultRepositoryImpl
                        .LoopResultChunk.class));
        verify(store, times(1)).compareAndSet(
                anyString(),
                eq(0L),
                any(Serializable.class),
                eq(1L),
                anyString(),
                eq(1L),
                anyString(),
                eq(1L),
                any(Duration.class));
    }

    /**
     * 验证调用方修改已提交的可变输出时不会污染下一轮活动分块写入。
     *
     * @throws Exception 测试依赖注入失败时抛出
     */
    @Test
    public void shouldIsolateMutableOutputValuesFromActiveCache()
            throws Exception {
        LoopResultRepositoryImpl repository =
                new LoopResultRepositoryImpl();
        VersionedObjectStore store =
                mock(VersionedObjectStore.class);
        Field field = LoopResultRepositoryImpl.class
                .getDeclaredField("versionedObjectStore");
        field.setAccessible(true);
        field.set(repository, store);
        when(store.load(
                anyString(),
                eq(LoopResultRepositoryImpl
                        .LoopResultChunk.class)))
                .thenReturn(null);
        when(store.createIfAbsent(
                anyString(),
                any(Serializable.class),
                anyLong(),
                anyString(),
                anyLong(),
                anyString(),
                anyLong(),
                any(Duration.class)))
                .thenReturn(true);
        when(store.compareAndSet(
                anyString(),
                anyLong(),
                any(Serializable.class),
                anyLong(),
                anyString(),
                anyLong(),
                anyString(),
                anyLong(),
                any(Duration.class)))
                .thenReturn(true);
        Map<String, Object> firstValue =
                new LinkedHashMap<>();
        firstValue.put("name", "original");

        repository.append(
                "instance",
                1L,
                "claim",
                1L,
                "instance:mutable-output",
                0,
                Map.of("value", firstValue));
        firstValue.put("name", "changed");
        repository.append(
                "instance",
                1L,
                "claim",
                1L,
                "instance:mutable-output",
                1,
                Map.of("value", Map.of(
                        "name", "second")));

        org.mockito.ArgumentCaptor<Serializable>
                chunkCaptor =
                org.mockito.ArgumentCaptor.forClass(
                        Serializable.class);
        verify(store).compareAndSet(
                anyString(),
                eq(0L),
                chunkCaptor.capture(),
                eq(1L),
                anyString(),
                eq(1L),
                anyString(),
                eq(1L),
                any(Duration.class));
        LoopResultRepositoryImpl.LoopResultChunk
                committed =
                (LoopResultRepositoryImpl.LoopResultChunk)
                        chunkCaptor.getValue();
        @SuppressWarnings("unchecked")
        Map<String, Object> committedFirst =
                (Map<String, Object>)
                        committed.getValues()
                                .get("value")
                                .get(0);
        Assert.assertEquals(
                "original",
                committedFirst.get("name"));
    }

    /**
     * 验证跨分块后活动输出缓存只保留当前分块。
     *
     * @throws Exception 测试依赖注入或反射读取失败时抛出
     */
    @Test
    public void shouldKeepOnlyCurrentOutputChunkInActiveCache()
            throws Exception {
        LoopResultRepositoryImpl repository =
                new LoopResultRepositoryImpl();
        VersionedObjectStore store =
                mock(VersionedObjectStore.class);
        Field storeField = LoopResultRepositoryImpl.class
                .getDeclaredField("versionedObjectStore");
        storeField.setAccessible(true);
        storeField.set(repository, store);
        when(store.createIfAbsent(
                anyString(),
                any(Serializable.class),
                anyLong(),
                anyString(),
                anyLong(),
                anyString(),
                anyLong(),
                any(Duration.class)))
                .thenReturn(true);

        repository.append(
                "instance",
                1L,
                "claim",
                1L,
                "instance:chunk-release",
                0,
                Map.of("value", "first"));
        repository.append(
                "instance",
                1L,
                "claim",
                1L,
                "instance:chunk-release",
                LoopResultRepositoryImpl.CHUNK_SIZE,
                Map.of("value", "next"));

        Field cacheField = LoopResultRepositoryImpl.class
                .getDeclaredField("outputChunkCache");
        cacheField.setAccessible(true);
        Object cache = cacheField.get(repository);
        Field valuesField = cache.getClass()
                .getDeclaredField("values");
        valuesField.setAccessible(true);
        @SuppressWarnings("unchecked")
        Map<String, Object> cachedValues =
                (Map<String, Object>)
                        valuesField.get(cache);
        Assert.assertEquals(
                1, cachedValues.size());
    }

    /**
     * 验证物化中失去 fencing 守卫后立即停止并清理本 owner 已写分块。
     *
     * @throws Exception 测试依赖注入失败时抛出
     */
    @Test
    public void shouldCleanupPartialInputWhenClaimIsLost()
            throws Exception {
        LoopResultRepositoryImpl repository =
                new LoopResultRepositoryImpl();
        VersionedObjectStore store =
                mock(VersionedObjectStore.class);
        Field field = LoopResultRepositoryImpl.class
                .getDeclaredField("versionedObjectStore");
        field.setAccessible(true);
        field.set(repository, store);
        when(store.createIfAbsent(
                anyString(),
                any(Serializable.class),
                anyLong(),
                anyString(),
                anyLong(),
                anyString(),
                anyLong(),
                any(Duration.class)))
                .thenReturn(true, false);
        List<Integer> input = java.util.stream.IntStream
                .range(0, LoopResultRepositoryImpl.CHUNK_SIZE + 1)
                .boxed()
                .toList();

        try {
            repository.storeInput(
                    "instance",
                    1L,
                    "claim",
                    1L,
                    "instance:guarded-input",
                    input,
                    10_000L);
            Assert.fail("lost claim must stop materialization");
        } catch (TriggerClaimLostException expected) {
            // 第二个分块守卫失败后立即退出。
        }

        verify(store).deleteAll(
                org.mockito.ArgumentMatchers.argThat(
                        keys -> keys.size() == 1));
    }

    /**
     * 验证锁外物化只依赖稳定 claim，合法实例锁代际推进不会中断后续分块。
     *
     * @throws Exception 测试依赖注入失败时抛出
     */
    @Test
    public void shouldMaterializeAllChunksWithStableClaimGuard()
            throws Exception {
        LoopResultRepositoryImpl repository =
                new LoopResultRepositoryImpl();
        VersionedObjectStore store =
                mock(VersionedObjectStore.class);
        Field field = LoopResultRepositoryImpl.class
                .getDeclaredField("versionedObjectStore");
        field.setAccessible(true);
        field.set(repository, store);
        when(store.createIfAbsent(
                anyString(),
                any(Serializable.class),
                anyLong(),
                anyString(),
                anyLong(),
                any(Duration.class))).thenReturn(true);
        List<Integer> input = java.util.stream.IntStream
                .range(0, LoopResultRepositoryImpl.CHUNK_SIZE + 1)
                .boxed()
                .toList();

        int stored = repository.storeProducedInput(
                "instance",
                0L,
                "claim",
                7L,
                "instance:stable-input",
                sink -> input.forEach(sink),
                10_000L);

        Assert.assertEquals(input.size(), stored);
        verify(store, times(3)).createIfAbsent(
                anyString(),
                any(Serializable.class),
                anyLong(),
                anyString(),
                eq(7L),
                any(Duration.class));
    }

    /**
     * 使用内存 Map 隔离 JetCache 的测试仓储。
     */
    private static final class InMemoryLoopResultRepository extends LoopResultRepositoryImpl {

        private final Map<String, Object> values = new LinkedHashMap<>();

        /**
         * 将分块写入测试内存。
         *
         * @param key 缓存键
         * @param value 缓存值
         */
        @Override
        protected void putCache(String key, Object value) {
            values.put(key, value);
        }

        /**
         * 从测试内存读取分块。
         *
         * @param key 缓存键
         * @param clazz 期望类型
         * @param <T> 缓存值类型
         * @return 命中的分块
         */
        @Override
        protected <T> T getCache(String key, Class<T> clazz) {
            return clazz.cast(values.get(key));
        }
    }
}
