package com.easyagents.flow.core.test;

import com.easyagents.flow.core.chain.repository.InMemoryLoopResultRepository;
import com.easyagents.flow.core.chain.repository.LoopInputReference;
import com.easyagents.flow.core.chain.repository.LoopResultReference;
import com.easyagents.flow.core.chain.runtime.ExecutionBudgetExceededException;
import org.junit.Assert;
import org.junit.Test;

import java.util.Arrays;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Iterator;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 循环结果引用批量还原回归测试。
 */
public class LoopResultReferenceResolverTest {

    /**
     * 验证同一循环的多个输出只触发一次仓储加载。
     */
    @Test
    @SuppressWarnings("unchecked")
    public void shouldBatchOutputsFromSameLoopResult() {
        CountingLoopResultRepository repository = new CountingLoopResultRepository();
        repository.append("result-1", 0, Map.of("a", "a0", "b", "b0"));
        repository.append("result-1", 1, Map.of("a", "a1", "b", "b1"));
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("a", new LoopResultReference("result-1", 2, "a"));
        value.put("nested", List.of(new LoopResultReference("result-1", 2, "b")));

        Map<String, Object> resolved =
                (Map<String, Object>) repository.resolveReferences(value);

        Assert.assertEquals(Arrays.asList("a0", "a1"), resolved.get("a"));
        Assert.assertEquals(
                Arrays.asList("b0", "b1"),
                ((List<Object>) resolved.get("nested")).get(0));
        Assert.assertEquals(1, repository.getLoadCount());
    }

    /**
     * 验证数组输出按轮次和轮内顺序只扁平一层。
     */
    @Test
    public void shouldFlattenArrayOutputOnceInStableOrder() {
        InMemoryLoopResultRepository repository =
                new InMemoryLoopResultRepository();
        repository.append(
                "result-flatten",
                0,
                Map.of(
                        "items",
                        List.of(
                                List.of("a"),
                                List.of("b"))));
        repository.append(
                "result-flatten",
                1,
                Map.of(
                        "items",
                        List.of(
                                List.of("c"))));

        Object resolved = repository.resolve(
                new LoopResultReference(
                        "result-flatten",
                        2,
                        "items",
                        true));

        Assert.assertEquals(
                List.of(
                        List.of("a"),
                        List.of("b"),
                        List.of("c")),
                resolved);
    }

    /**
     * 验证启用扁平聚合后拒绝某轮标量结果。
     */
    @Test
    public void shouldRejectScalarIterationWhenFlattening() {
        InMemoryLoopResultRepository repository =
                new InMemoryLoopResultRepository();
        repository.append(
                "result-invalid",
                0,
                Map.of("items", List.of("a")));
        repository.append(
                "result-invalid",
                1,
                Map.of("items", "scalar"));

        try {
            repository.resolve(
                    new LoopResultReference(
                            "result-invalid",
                            2,
                            "items",
                            true));
            Assert.fail("scalar iteration should be rejected");
        } catch (IllegalStateException expected) {
            Assert.assertTrue(
                    expected.getMessage().contains("items"));
            Assert.assertTrue(
                    expected.getMessage().contains("iteration 1"));
        }
    }

    /**
     * 验证循环完成时只给指定输出携带扁平聚合策略。
     */
    @Test
    public void shouldCreateReferencesWithPerOutputAggregation() {
        InMemoryLoopResultRepository repository =
                new InMemoryLoopResultRepository();

        Map<String, Object> references = repository.references(
                "result-references",
                2,
                List.of("grouped", "flattened"),
                Set.of("flattened"));

        Assert.assertFalse(
                ((LoopResultReference) references.get("grouped"))
                        .isFlattenAggregation());
        Assert.assertTrue(
                ((LoopResultReference) references.get("flattened"))
                        .isFlattenAggregation());
    }

    /**
     * 验证无限 Iterable 在达到预算后立即停止物化。
     */
    @Test
    public void shouldStopInfiniteInputAtIterationBudget() {
        InMemoryLoopResultRepository repository = new InMemoryLoopResultRepository();
        AtomicInteger reads = new AtomicInteger();
        Iterable<Integer> infinite = () -> new Iterator<Integer>() {
            @Override
            public boolean hasNext() {
                return true;
            }

            @Override
            public Integer next() {
                return reads.incrementAndGet();
            }
        };

        try {
            repository.storeInput("bounded-input", infinite, 3L);
            Assert.fail("infinite input should exceed the iteration budget");
        } catch (ExecutionBudgetExceededException expected) {
            Assert.assertTrue(expected.getMessage().contains("more than 3"));
        }

        Assert.assertEquals(3, reads.get());
    }

    /**
     * 验证外置循环输入在业务读取边界仍还原为原顺序列表。
     */
    @Test
    public void shouldResolveExternalizedLoopInput() {
        InMemoryLoopResultRepository repository =
                new InMemoryLoopResultRepository();
        String resultId = "input-reference";
        Assert.assertEquals(
                3,
                repository.storeInput(
                        resultId,
                        List.of("first", "second", "third")));

        Object resolved = repository.resolveReferences(
                new LoopInputReference(resultId, 3));

        Assert.assertEquals(
                List.of("first", "second", "third"),
                resolved);
    }

    /**
     * 记录加载次数的循环结果仓储。
     */
    private static final class CountingLoopResultRepository
            extends InMemoryLoopResultRepository {

        private final AtomicInteger loadCount = new AtomicInteger();

        /**
         * {@inheritDoc}
         */
        @Override
        public Map<String, Object> load(
                String resultId, int iterationCount, List<String> outputNames) {
            loadCount.incrementAndGet();
            return super.load(resultId, iterationCount, outputNames);
        }

        /**
         * 获取仓储加载次数。
         *
         * @return 加载次数
         */
        private int getLoadCount() {
            return loadCount.get();
        }
    }
}
