package com.easyagents.agent.runtime.agentscope;

import com.easyagents.agent.runtime.AgentDefinition;
import com.easyagents.agent.runtime.AgentRuntimeExecutionContext;
import com.easyagents.agent.runtime.tool.AgentToolInvoker;
import com.easyagents.agent.runtime.tool.AgentToolResult;
import com.easyagents.agent.runtime.tool.AgentToolSpec;
import io.agentscope.core.message.ToolResultBlock;
import io.agentscope.core.message.ToolUseBlock;
import io.agentscope.core.tool.AgentTool;
import io.agentscope.core.tool.ToolCallParam;
import org.junit.Assert;
import org.junit.Test;

import java.lang.reflect.Proxy;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;

/**
 * {@link AgentScopeToolAdapter} 回归测试。
 */
public class AgentScopeToolAdapterTest {

    /**
     * 验证 Reactor 工作线程执行工具时使用调用器所属类加载器，并在结束后恢复原上下文。
     */
    @Test
    public void shouldUseInvokerClassLoaderAndRestoreWorkerContext() {
        ClassLoader parentClassLoader = AgentToolInvoker.class.getClassLoader();
        ClassLoader invocationClassLoader = new ClassLoader(parentClassLoader) {
        };
        AtomicReference<ClassLoader> observedClassLoader = new AtomicReference<>();
        AgentToolInvoker invoker = (AgentToolInvoker) Proxy.newProxyInstance(
                invocationClassLoader,
                new Class<?>[]{AgentToolInvoker.class},
                (proxy, method, arguments) -> {
                    observedClassLoader.set(Thread.currentThread().getContextClassLoader());
                    return AgentToolResult.success("ok");
                });
        AgentTool tool = new AgentScopeToolAdapter().adapt(toolSpec(), invoker, executionContext());
        Thread currentThread = Thread.currentThread();
        ClassLoader originalClassLoader = currentThread.getContextClassLoader();
        ClassLoader workerClassLoader = new ClassLoader(originalClassLoader) {
        };

        try {
            currentThread.setContextClassLoader(workerClassLoader);
            ToolResultBlock result = tool.callAsync(toolCall()).block();

            Assert.assertNotNull(result);
            Assert.assertSame(invocationClassLoader, observedClassLoader.get());
            Assert.assertSame(workerClassLoader, currentThread.getContextClassLoader());
        } finally {
            currentThread.setContextClassLoader(originalClassLoader);
        }
    }

    /**
     * 创建测试工具声明。
     *
     * @return 工具声明
     */
    private AgentToolSpec toolSpec() {
        AgentToolSpec toolSpec = new AgentToolSpec();
        toolSpec.setName("test_tool");
        toolSpec.setDescription("Test tool");
        toolSpec.setParametersSchema(Map.of(
                "type", "object",
                "properties", Map.of(),
                "additionalProperties", false));
        return toolSpec;
    }

    /**
     * 创建测试运行上下文。
     *
     * @return 运行上下文
     */
    private AgentRuntimeExecutionContext executionContext() {
        AgentDefinition definition = new AgentDefinition();
        definition.setAgentId("agent-1");
        AgentRuntimeExecutionContext context = new AgentRuntimeExecutionContext();
        context.setAgentDefinition(definition);
        context.setRequestId("request-1");
        context.setTraceId("trace-1");
        context.setSessionId("session-1");
        return context;
    }

    /**
     * 创建测试工具调用参数。
     *
     * @return 工具调用参数
     */
    private ToolCallParam toolCall() {
        ToolUseBlock toolUseBlock = ToolUseBlock.builder()
                .id("call-1")
                .name("test_tool")
                .input(Map.of())
                .content("{}")
                .build();
        return ToolCallParam.builder()
                .toolUseBlock(toolUseBlock)
                .input(Map.of())
                .build();
    }
}
