package com.easyagents.flow.core.test;

import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainDefinition;
import com.easyagents.flow.core.chain.ChainState;
import com.easyagents.flow.core.chain.Parameter;
import com.easyagents.flow.core.chain.RefType;
import com.easyagents.flow.core.chain.repository.InMemoryLoopResultRepository;
import com.easyagents.flow.core.chain.repository.LoopInputReference;
import com.easyagents.flow.core.node.StartNode;
import org.junit.Assert;
import org.junit.Test;

import java.lang.reflect.Field;
import java.util.Collections;
import java.util.List;
import java.util.Map;

/**
 * 验证固定值模板在运行时可以直接读取链路 memory 中的节点作用域变量。
 */
public class ChainTemplateContextTest {

    @Test
    public void shouldResolveScopedNumericMemoryKeyInFixedParameterTemplate() {
        ChainState state = new ChainState();
        state.getMemory().put("node_1.123", "7");

        Parameter parameter = new Parameter();
        parameter.setName("nextValue");
        parameter.setRefType(RefType.FIXED);
        parameter.setValue("{{node_1.123}}");

        StartNode node = new StartNode();
        node.setName("开始节点");

        Map<String, Object> result = state.resolveParameters(node, Collections.singletonList(parameter));

        Assert.assertEquals("7", result.get("nextValue"));
    }

    /**
     * 验证审计参数快照保留直接大型引用，不在工作流线程读取完整输入。
     *
     * @throws Exception 线程上下文反射失败时抛出
     */
    @Test
    public void shouldPreserveDirectReferenceForAuditWithoutLoadingInput()
            throws Exception {
        CountingLoopRepository repository =
                new CountingLoopRepository();
        Chain chain = new Chain(
                new ChainDefinition(),
                "instance-audit");
        chain.setLoopResultRepository(repository);
        Field currentChain = Chain.class
                .getDeclaredField(
                        "EXECUTION_THREAD_LOCAL");
        currentChain.setAccessible(true);
        @SuppressWarnings("unchecked")
        ThreadLocal<Chain> context =
                (ThreadLocal<Chain>) currentChain.get(null);
        LoopInputReference reference =
                new LoopInputReference(
                        "instance-audit:dataset",
                        100_000);
        ChainState state = new ChainState();
        state.getMemory().put(
                "dataset.data", reference);
        Parameter parameter = new Parameter();
        parameter.setName("items");
        parameter.setRefType(RefType.REF);
        parameter.setRef("dataset.data");
        StartNode node = new StartNode();
        node.setParameters(List.of(parameter));

        context.set(chain);
        try {
            Map<String, Object> result =
                    state.resolveParametersPreservingReferences(
                            node);

            Assert.assertSame(
                    reference, result.get("items"));
            Assert.assertEquals(
                    0, repository.loadInputCalls);
        } finally {
            context.remove();
        }
    }

    /**
     * 验证无关固定参数不会触发 memory 中大型引用的物化。
     *
     * @throws Exception 线程上下文反射失败时抛出
     */
    @Test
    public void shouldNotLoadUnrelatedReferenceForFixedAuditParameter()
            throws Exception {
        CountingLoopRepository repository =
                new CountingLoopRepository();
        Chain chain = new Chain(
                new ChainDefinition(),
                "instance-fixed-audit");
        chain.setLoopResultRepository(repository);
        Field currentChain = Chain.class
                .getDeclaredField(
                        "EXECUTION_THREAD_LOCAL");
        currentChain.setAccessible(true);
        @SuppressWarnings("unchecked")
        ThreadLocal<Chain> context =
                (ThreadLocal<Chain>) currentChain.get(null);
        ChainState state = new ChainState();
        state.getMemory().put(
                "dataset.data",
                new LoopInputReference(
                        "instance-fixed-audit:dataset",
                        100_000));
        state.getMemory().put(
                "small.value", "ok");
        Parameter parameter = new Parameter();
        parameter.setName("description");
        parameter.setRefType(RefType.FIXED);
        parameter.setValue("{{small.value}}");
        StartNode node = new StartNode();
        node.setParameters(List.of(parameter));

        context.set(chain);
        try {
            Map<String, Object> result =
                    state.resolveParametersPreservingReferences(
                            node);

            Assert.assertEquals(
                    "ok",
                    result.get("description"));
            Assert.assertEquals(
                    0, repository.loadInputCalls);
        } finally {
            context.remove();
        }
    }

    /**
     * 验证普通节点固定参数同样只还原模板实际读取的大型引用。
     *
     * @throws Exception 线程上下文反射失败时抛出
     */
    @Test
    public void shouldNotLoadUnrelatedReferenceForNormalFixedParameter()
            throws Exception {
        CountingLoopRepository repository =
                new CountingLoopRepository();
        Chain chain = new Chain(
                new ChainDefinition(),
                "instance-fixed-normal");
        chain.setLoopResultRepository(repository);
        Field currentChain = Chain.class
                .getDeclaredField(
                        "EXECUTION_THREAD_LOCAL");
        currentChain.setAccessible(true);
        @SuppressWarnings("unchecked")
        ThreadLocal<Chain> context =
                (ThreadLocal<Chain>) currentChain.get(null);
        ChainState state = new ChainState();
        state.getMemory().put(
                "dataset.data",
                new LoopInputReference(
                        "instance-fixed-normal:dataset",
                        100_000));
        state.getMemory().put(
                "small.value", "ok");
        Parameter parameter = new Parameter();
        parameter.setName("description");
        parameter.setRefType(RefType.FIXED);
        parameter.setValue("{{small.value}}");
        StartNode node = new StartNode();
        node.setParameters(List.of(parameter));

        context.set(chain);
        try {
            Map<String, Object> result =
                    state.resolveParameters(node);

            Assert.assertEquals(
                    "ok",
                    result.get("description"));
            Assert.assertEquals(
                    0, repository.loadInputCalls);
        } finally {
            context.remove();
        }
    }

    /**
     * 验证同一模板重复读取一个大型引用时只执行一次仓储还原。
     *
     * @throws Exception 线程上下文反射失败时抛出
     */
    @Test
    public void shouldMemoizeRepeatedReferenceWithinOneTemplate()
            throws Exception {
        CountingLoopRepository repository =
                new CountingLoopRepository();
        String resultId =
                "instance-fixed-repeat:dataset";
        repository.storeInput(
                resultId, List.of("item"));
        Chain chain = new Chain(
                new ChainDefinition(),
                "instance-fixed-repeat");
        chain.setLoopResultRepository(repository);
        Field currentChain = Chain.class
                .getDeclaredField(
                        "EXECUTION_THREAD_LOCAL");
        currentChain.setAccessible(true);
        @SuppressWarnings("unchecked")
        ThreadLocal<Chain> context =
                (ThreadLocal<Chain>) currentChain.get(null);
        ChainState state = new ChainState();
        state.getMemory().put(
                "dataset.data",
                new LoopInputReference(
                        resultId, 1));
        Parameter parameter = new Parameter();
        parameter.setName("description");
        parameter.setRefType(RefType.FIXED);
        parameter.setValue(
                "{{dataset.data}}/{{dataset.data}}");
        StartNode node = new StartNode();
        node.setParameters(List.of(parameter));

        context.set(chain);
        try {
            Map<String, Object> result =
                    state.resolveParameters(node);

            Assert.assertEquals(
                    "[item]/[item]",
                    result.get("description"));
            Assert.assertEquals(
                    1, repository.loadInputCalls);
        } finally {
            context.remove();
        }
    }

    /**
     * 记录大型输入加载次数的仓储。
     */
    private static final class CountingLoopRepository
            extends InMemoryLoopResultRepository {

        private int loadInputCalls;

        /**
         * {@inheritDoc}
         */
        @Override
        public List<Object> loadInput(
                LoopInputReference reference) {
            loadInputCalls++;
            return super.loadInput(reference);
        }
    }
}
