package tech.easyflow.agent.runtime;

import com.easyagents.agent.runtime.agentscope.AgentScopeMessageAdapter;
import com.easyagents.agent.runtime.model.AgentGenerationOptions;
import com.easyagents.agent.runtime.model.AgentModelFactory;
import com.easyagents.agent.runtime.model.AgentModelSpec;
import io.agentscope.core.message.Base64Source;
import io.agentscope.core.message.ImageBlock;
import io.agentscope.core.message.Msg;
import io.agentscope.core.message.MsgRole;
import io.agentscope.core.message.TextBlock;
import io.agentscope.core.message.ToolUseBlock;
import io.agentscope.core.model.ChatResponse;
import io.agentscope.core.model.GenerateOptions;
import io.agentscope.core.model.ToolChoice;
import io.agentscope.core.model.ToolSchema;
import org.junit.Assert;
import org.junit.Test;
import reactor.core.publisher.Flux;
import tech.easyflow.ai.entity.Model;
import tech.easyflow.ai.entity.ModelProvider;
import tech.easyflow.ai.service.support.VlmVerificationImage;
import tech.easyflow.ai.service.verification.ChatModelVerificationResult;
import tech.easyflow.ai.service.verification.ModelVerificationStatus;
import tech.easyflow.common.web.exceptions.BusinessException;

import java.math.BigInteger;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.Map;

/**
 * AgentScope 单次优先模型连接与工具能力验证测试。
 */
public class AgentScopeChatModelConnectivityVerifierTest {

    private static final String TEST_NONCE = "probe-nonce";

    /**
     * 验证一次请求正确返回工具调用时同时确认连接和工具能力。
     */
    @Test
    public void shouldPassConnectionAndToolProbeInOneRequest() {
        RecordingModelFactory factory = new RecordingModelFactory(
                Flux.just(toolResponse(TEST_NONCE, null)));

        ChatModelVerificationResult result = verifier(factory).verify(model(false));

        Assert.assertEquals(ModelVerificationStatus.PASSED, result.getStatus());
        Assert.assertEquals(ModelVerificationStatus.PASSED, result.getNonStreaming());
        Assert.assertEquals(ModelVerificationStatus.SKIPPED, result.getStreaming());
        Assert.assertEquals(Boolean.TRUE, result.getSupportTool());
        Assert.assertEquals("验证通过", result.getMessage());
        Assert.assertEquals(List.of(false), factory.getFactoryStreams());
        Assert.assertEquals(List.of(false), factory.getRequestStreams());
        Assert.assertEquals(List.of(1), factory.getToolCounts());
        Assert.assertTrue(factory.getToolChoices().get(0) instanceof ToolChoice.Specific);
        Assert.assertEquals(MsgRole.SYSTEM, factory.getMessageBatches().get(0).get(0).getRole());
        Assert.assertEquals(MsgRole.USER, factory.getMessageBatches().get(0).get(1).getRole());
    }

    /**
     * 验证能力探测达到输出上限时使用更大预算重试一次。
     */
    @Test
    public void shouldRetryProbeWithLargerBudgetWhenResponseIsLengthLimited() {
        RecordingModelFactory factory = new RecordingModelFactory(
                Flux.just(lengthLimitedResponse()),
                Flux.just(toolResponse(TEST_NONCE, null)));

        ChatModelVerificationResult result = verifier(factory).verify(model(false));

        Assert.assertEquals(ModelVerificationStatus.PASSED, result.getStatus());
        Assert.assertEquals(Boolean.TRUE, result.getSupportTool());
        Assert.assertEquals(List.of(1, 1), factory.getToolCounts());
        Assert.assertEquals(List.of(128, 512), factory.getFactoryMaxTokens());
        Assert.assertEquals(List.of(128, 512), factory.getRequestMaxTokens());
    }

    /**
     * 验证自部署 vLLM/SGLang 入口通过聊天模板参数关闭思考。
     */
    @Test
    public void shouldDisableThinkingThroughChatTemplateKwargsForSelfHostedEndpoint() {
        RecordingModelFactory factory = new RecordingModelFactory(
                Flux.just(toolResponse(TEST_NONCE, null)));

        ChatModelVerificationResult result = verifier(factory).verify(
                model(false, Map.of(), "self-hosted"));

        Assert.assertEquals(ModelVerificationStatus.PASSED, result.getStatus());
        Map<String, Object> bodyParams = factory.getFactoryAdditionalBodyParams().get(0);
        Assert.assertEquals(Boolean.FALSE, bodyParams.get("enable_thinking"));
        Assert.assertEquals(
                Map.of("enable_thinking", false),
                bodyParams.get("chat_template_kwargs"));
    }

    /**
     * 验证模型返回普通文本时连接通过但工具能力保持关闭。
     */
    @Test
    public void shouldPassConnectionAndMarkToolUnsupportedWhenTextReturned() {
        RecordingModelFactory factory = new RecordingModelFactory(
                Flux.just(textResponse("你好")));

        ChatModelVerificationResult result = verifier(factory).verify(model(false));

        Assert.assertEquals(ModelVerificationStatus.PASSED, result.getStatus());
        Assert.assertEquals(Boolean.FALSE, result.getSupportTool());
        Assert.assertEquals(List.of(1), factory.getToolCounts());
    }

    /**
     * 验证接口明确拒绝工具参数时仅追加一次普通连接兜底。
     */
    @Test
    public void shouldFallbackToPlainRequestWhenToolChoiceIsRejected() {
        RecordingModelFactory factory = new RecordingModelFactory(
                Flux.error(new IllegalArgumentException("tool_choice is unsupported")),
                Flux.just(textResponse("你好")));

        ChatModelVerificationResult result = verifier(factory).verify(model(false));

        Assert.assertEquals(ModelVerificationStatus.PASSED, result.getStatus());
        Assert.assertEquals(Boolean.FALSE, result.getSupportTool());
        Assert.assertEquals(List.of(1, 0), factory.getToolCounts());
        Assert.assertTrue(factory.getToolChoices().get(0) instanceof ToolChoice.Specific);
        Assert.assertNull(factory.getToolChoices().get(1));
    }

    /**
     * 验证普通连接异常不会被工具兼容兜底掩盖。
     */
    @Test
    public void shouldFailWhenConnectivityRequestFails() {
        RecordingModelFactory factory = new RecordingModelFactory(
                Flux.error(new IllegalStateException("connection failed")));

        try {
            verifier(factory).verify(model(false));
            Assert.fail("Expected connectivity verification failure");
        } catch (BusinessException exception) {
            Assert.assertTrue(exception.getMessage().contains("连接验证失败"));
            Assert.assertFalse(exception.getMessage().contains("connection failed"));
        }
        Assert.assertEquals(List.of(1), factory.getToolCounts());
    }

    /**
     * 验证 VLM 在同一次工具调用中返回图片验证码。
     */
    @Test
    public void shouldVerifyVlmAndToolCallInOneRequest() {
        RecordingModelFactory factory = new RecordingModelFactory(
                Flux.just(toolResponse(TEST_NONCE, VlmVerificationImage.VERIFICATION_CODE)));

        ChatModelVerificationResult result = verifier(factory).verify(model(true));

        Assert.assertEquals(ModelVerificationStatus.PASSED, result.getStatus());
        Assert.assertEquals(Boolean.TRUE, result.getSupportTool());
        Msg message = factory.getMessageBatches().get(0).get(1);
        ImageBlock image = message.getContent().stream()
                .filter(ImageBlock.class::isInstance)
                .map(ImageBlock.class::cast)
                .findFirst()
                .orElseThrow();
        Assert.assertTrue(image.getSource() instanceof Base64Source);
        byte[] imageBytes = Base64.getDecoder().decode(((Base64Source) image.getSource()).getData());
        Assert.assertEquals((byte) 0x89, imageBytes[0]);
        Assert.assertEquals((byte) 0x50, imageBytes[1]);
    }

    /**
     * 验证严格内容块模式会先发送包含图片首轮和文本追问的完整上下文。
     */
    @Test
    public void shouldVerifyStrictMessageContentFormatWithVlmHistory() {
        RecordingModelFactory factory = new RecordingModelFactory(
                Flux.just(textResponse("你好")),
                Flux.just(toolResponse(TEST_NONCE, VlmVerificationImage.VERIFICATION_CODE)));

        ChatModelVerificationResult result = verifier(factory).verify(model(
                true,
                Map.of("agentMessageContentFormat", "TEXT_PARTS")));

        Assert.assertEquals(ModelVerificationStatus.PASSED, result.getStatus());
        Assert.assertEquals(List.of(0, 1), factory.getToolCounts());
        List<Msg> history = factory.getMessageBatches().get(0);
        Assert.assertEquals(List.of(
                        MsgRole.SYSTEM,
                        MsgRole.USER,
                        MsgRole.ASSISTANT,
                        MsgRole.USER),
                history.stream().map(Msg::getRole).toList());
        Assert.assertTrue(history.get(1).getContent().stream().anyMatch(ImageBlock.class::isInstance));
        Assert.assertEquals("请直接回复“你好”，不要补充其他内容。",
                history.get(3).getTextContent());
    }

    /**
     * 创建待验证模型。
     *
     * @param supportImage 是否支持图片
     * @return 测试模型
     */
    private Model model(boolean supportImage) {
        return model(supportImage, Map.of());
    }

    /**
     * 创建携带指定 options 的待验证模型。
     *
     * @param supportImage 是否支持图片
     * @param options 模型扩展配置
     * @return 测试模型
     */
    private Model model(boolean supportImage, Map<String, Object> options) {
        return model(supportImage, options, "gpustack");
    }

    /**
     * 创建携带指定供应商类型的待验证模型。
     *
     * @param supportImage 是否支持图片
     * @param options 模型扩展配置
     * @param providerType 供应商类型
     * @return 测试模型
     */
    private Model model(boolean supportImage,
                        Map<String, Object> options,
                        String providerType) {
        Model model = new Model();
        model.setId(BigInteger.TEN);
        model.setModelName("test-model");
        model.setEndpoint("http://model.example.com");
        model.setRequestPath("/v1/chat/completions");
        model.setApiKey("test-key");
        model.setSupportImage(supportImage);
        model.setOptions(options);
        ModelProvider provider = new ModelProvider();
        provider.setProviderType(providerType);
        model.setModelProvider(provider);
        return model;
    }

    /**
     * 创建使用固定随机值的验证器。
     *
     * @param factory 记录型模型工厂
     * @return 验证器
     */
    private AgentScopeChatModelConnectivityVerifier verifier(RecordingModelFactory factory) {
        return new AgentScopeChatModelConnectivityVerifier(
                factory,
                new AgentScopeMessageAdapter(),
                Duration.ofSeconds(2),
                () -> TEST_NONCE);
    }

    /**
     * 创建文本响应。
     *
     * @param text 文本内容
     * @return AgentScope 响应
     */
    private ChatResponse textResponse(String text) {
        return ChatResponse.builder()
                .content(List.of(TextBlock.builder().text(text).build()))
                .build();
    }

    /**
     * 创建达到输出上限的响应。
     *
     * @return 输出被截断的 AgentScope 响应
     */
    private ChatResponse lengthLimitedResponse() {
        return ChatResponse.builder()
                .content(List.of(TextBlock.builder().text("incomplete").build()))
                .finishReason("length")
                .build();
    }

    /**
     * 创建探测工具调用响应。
     *
     * @param nonce 随机校验值
     * @param imageCode 图片验证码
     * @return AgentScope 响应
     */
    private ChatResponse toolResponse(String nonce, String imageCode) {
        Map<String, Object> input = imageCode == null
                ? Map.of("nonce", nonce)
                : Map.of("nonce", nonce, "imageCode", imageCode);
        return ChatResponse.builder()
                .content(List.of(ToolUseBlock.builder()
                        .id("call-probe")
                        .name("easyflow_capability_probe")
                        .input(input)
                        .build()))
                .build();
    }

    /**
     * 按调用顺序返回预设响应并记录真实请求参数的模型工厂。
     */
    private static final class RecordingModelFactory
            implements AgentModelFactory<io.agentscope.core.model.Model> {

        /** 每次调用的预设响应。 */
        private final List<Flux<ChatResponse>> responses;
        /** 模型工厂收到的流式参数。 */
        private final List<Boolean> factoryStreams = new ArrayList<>();
        /** 模型工厂收到的最大输出 Token 数。 */
        private final List<Integer> factoryMaxTokens = new ArrayList<>();
        /** 模型工厂收到的额外请求体参数。 */
        private final List<Map<String, Object>> factoryAdditionalBodyParams = new ArrayList<>();
        /** 模型请求收到的流式参数。 */
        private final List<Boolean> requestStreams = new ArrayList<>();
        /** 模型请求收到的最大输出 Token 数。 */
        private final List<Integer> requestMaxTokens = new ArrayList<>();
        /** 模型请求收到的消息批次。 */
        private final List<List<Msg>> messageBatches = new ArrayList<>();
        /** 每次请求携带的工具数量。 */
        private final List<Integer> toolCounts = new ArrayList<>();
        /** 每次请求使用的工具选择策略。 */
        private final List<ToolChoice> toolChoices = new ArrayList<>();

        /**
         * 创建记录型模型工厂。
         *
         * @param responses 每次调用的预设响应
         */
        @SafeVarargs
        private RecordingModelFactory(Flux<ChatResponse>... responses) {
            this.responses = List.of(responses);
        }

        /**
         * 创建当前验证请求使用的模型。
         *
         * @param modelSpec 模型声明
         * @param generationOptions 生成参数
         * @return 测试模型
         */
        @Override
        public io.agentscope.core.model.Model create(
                AgentModelSpec modelSpec,
                AgentGenerationOptions generationOptions) {
            int requestIndex = factoryStreams.size();
            factoryStreams.add(Boolean.TRUE.equals(generationOptions.getStream()));
            factoryMaxTokens.add(generationOptions.getMaxTokens());
            factoryAdditionalBodyParams.add(Map.copyOf(
                    generationOptions.getAdditionalBodyParams()));
            Flux<ChatResponse> response = responses.get(requestIndex);
            return new io.agentscope.core.model.Model() {
                /**
                 * 返回预设响应并记录请求参数。
                 *
                 * @param inputMessages 模型消息
                 * @param tools 工具声明
                 * @param options 生成参数
                 * @return 预设响应
                 */
                @Override
                public Flux<ChatResponse> stream(
                        List<Msg> inputMessages,
                        List<ToolSchema> tools,
                        GenerateOptions options) {
                    requestStreams.add(Boolean.TRUE.equals(options.getStream()));
                    requestMaxTokens.add(options.getMaxTokens());
                    messageBatches.add(List.copyOf(inputMessages));
                    toolCounts.add(tools.size());
                    toolChoices.add(options.getToolChoice());
                    return response;
                }

                /**
                 * 返回测试模型名称。
                 *
                 * @return 测试模型名称
                 */
                @Override
                public String getModelName() {
                    return modelSpec.getModelName();
                }
            };
        }

        /**
         * 获取模型工厂流式参数。
         *
         * @return 流式参数列表
         */
        private List<Boolean> getFactoryStreams() {
            return factoryStreams;
        }

        /**
         * 获取模型工厂收到的最大输出 Token 数。
         *
         * @return 最大输出 Token 数列表
         */
        private List<Integer> getFactoryMaxTokens() {
            return factoryMaxTokens;
        }

        /**
         * 获取模型工厂收到的额外请求体参数。
         *
         * @return 额外请求体参数列表
         */
        private List<Map<String, Object>> getFactoryAdditionalBodyParams() {
            return factoryAdditionalBodyParams;
        }

        /**
         * 获取请求流式参数。
         *
         * @return 流式参数列表
         */
        private List<Boolean> getRequestStreams() {
            return requestStreams;
        }

        /**
         * 获取模型请求收到的最大输出 Token 数。
         *
         * @return 最大输出 Token 数列表
         */
        private List<Integer> getRequestMaxTokens() {
            return requestMaxTokens;
        }

        /**
         * 获取每次请求的验证消息批次。
         *
         * @return 验证消息批次
         */
        private List<List<Msg>> getMessageBatches() {
            return messageBatches;
        }

        /**
         * 获取工具数量。
         *
         * @return 工具数量列表
         */
        private List<Integer> getToolCounts() {
            return toolCounts;
        }

        /**
         * 获取工具选择策略。
         *
         * @return 工具选择策略列表
         */
        private List<ToolChoice> getToolChoices() {
            return toolChoices;
        }
    }
}
