package tech.easyflow.ai.easyagentsflow.code;

import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainDefinition;
import com.easyagents.flow.core.chain.repository.InMemoryChainStateRepository;
import com.easyagents.flow.core.node.CodeNode;
import org.junit.Assert;
import org.junit.Assume;
import org.junit.Test;

import java.util.Map;
import java.util.UUID;

public class PythonRuntimeEngineTest {

    private static final String PYTHON_COMMAND = "python3";

    @Test
    public void testExecuteSuccess() {
        Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());

        Chain chain = createChain();
        CodeNode node = (CodeNode) chain.getDefinition().getNodeById("code-test");
        chain.getState().getMemory().put("a", 1);
        chain.getState().getMemory().put("b", 2);

        PythonRuntimeEngine engine = new PythonRuntimeEngine(PYTHON_COMMAND, 3000L, 65536, System.getProperty("java.io.tmpdir"));
        Map<String, Object> result = engine.execute("_result['sum'] = a + b\n_result['env_type'] = type(_env).__name__", node, chain);

        Assert.assertEquals(3, ((Number) result.get("sum")).intValue());
        Assert.assertEquals("dict", result.get("env_type"));
    }

    @Test
    public void testSyntaxError() {
        Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());

        Chain chain = createChain();
        CodeNode node = (CodeNode) chain.getDefinition().getNodeById("code-test");
        PythonRuntimeEngine engine = new PythonRuntimeEngine(PYTHON_COMMAND, 3000L, 65536, System.getProperty("java.io.tmpdir"));

        assertExecuteFail(engine, node, chain, "if True print('broken')", "执行失败");
    }

    @Test
    public void testTimeout() {
        Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());

        Chain chain = createChain();
        CodeNode node = (CodeNode) chain.getDefinition().getNodeById("code-test");
        PythonRuntimeEngine engine = new PythonRuntimeEngine(PYTHON_COMMAND, 100L, 65536, System.getProperty("java.io.tmpdir"));

        assertExecuteFail(engine, node, chain, "import time\ntime.sleep(1)\n_result['ok'] = 1", "超时");
    }

    @Test
    public void testOutputTooLarge() {
        Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());

        Chain chain = createChain();
        CodeNode node = (CodeNode) chain.getDefinition().getNodeById("code-test");
        PythonRuntimeEngine engine = new PythonRuntimeEngine(PYTHON_COMMAND, 3000L, 64, System.getProperty("java.io.tmpdir"));

        assertExecuteFail(engine, node, chain, "_result['text'] = 'a' * 2048", "输出超出限制");
    }

    @Test
    public void testCommandNotFound() {
        Chain chain = createChain();
        CodeNode node = (CodeNode) chain.getDefinition().getNodeById("code-test");
        PythonRuntimeEngine engine = new PythonRuntimeEngine("python3_not_exists_for_test", 1000L, 65536, System.getProperty("java.io.tmpdir"));

        assertExecuteFail(engine, node, chain, "_result['ok'] = 1", "执行失败");
    }

    private Chain createChain() {
        ChainDefinition definition = new ChainDefinition();
        CodeNode node = new CodeNode();
        node.setId("code-test");
        node.setEngine("python");
        definition.addNode(node);

        Chain chain = new Chain(definition, UUID.randomUUID().toString());
        chain.setChainStateRepository(new InMemoryChainStateRepository());
        return chain;
    }

    private void assertExecuteFail(PythonRuntimeEngine engine, CodeNode node, Chain chain, String code, String messageContains) {
        try {
            engine.execute(code, node, chain);
            Assert.fail("expected execute fail");
        } catch (RuntimeException e) {
            Assert.assertTrue(e.getMessage().contains(messageContains));
        }
    }
}
