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.TextBlock;
import io.agentscope.core.model.ChatResponse;
import io.agentscope.core.model.GenerateOptions;
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;

/**
 * AgentScope 双阶段模型连通性验证测试。
 */
public class AgentScopeChatModelConnectivityVerifierTest {

    /**
     * 验证非流式与流式阶段按顺序执行并返回通过状态。
     */
    @Test
    public void shouldPassWhenBothPhasesReturnText() {
        RecordingModelFactory factory = new RecordingModelFactory(
                Flux.just(response("你好")),
                Flux.just(response("你"), response("好")));
        AgentScopeChatModelConnectivityVerifier verifier = verifier(factory);

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

        Assert.assertEquals(ModelVerificationStatus.PASSED, result.getStatus());
        Assert.assertEquals("HTTP_1_1", result.getEffectiveHttpVersion());
        Assert.assertEquals(List.of(false, true), factory.getFactoryStreams());
        Assert.assertEquals(List.of(false, true), factory.getRequestStreams());
        Assert.assertEquals(List.of(false, false), factory.getEnableThinkingValues());
        Assert.assertEquals(List.of(false, false), factory.getChatTemplateThinkingValues());
    }

    /**
     * 验证基础连接通过而流式阶段失败时返回部分通过结果。
     */
    @Test
    public void shouldReturnPartialWhenStreamingPhaseFails() {
        RecordingModelFactory factory = new RecordingModelFactory(
                Flux.just(response("你好")),
                Flux.error(new IllegalStateException("stream failed")));

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

        Assert.assertEquals(ModelVerificationStatus.PARTIAL, result.getStatus());
        Assert.assertEquals(ModelVerificationStatus.PASSED, result.getNonStreaming());
        Assert.assertEquals(ModelVerificationStatus.FAILED, result.getStreaming());
        Assert.assertTrue(result.getMessage().contains("流式响应不可用"));
    }

    /**
     * 验证基础连接失败时立即终止且返回业务失败。
     */
    @Test
    public void shouldStopWhenNonStreamingPhaseFails() {
        RecordingModelFactory factory = new RecordingModelFactory(
                Flux.error(new IllegalStateException("connection failed")));

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

    /**
     * 验证流式阶段超时被归类为部分可用且不暴露底层异常。
     */
    @Test
    public void shouldReturnPartialWhenStreamingPhaseTimesOut() {
        RecordingModelFactory factory = new RecordingModelFactory(
                Flux.just(response("你好")),
                Flux.never());
        AgentScopeChatModelConnectivityVerifier verifier = new AgentScopeChatModelConnectivityVerifier(
                factory,
                new AgentScopeMessageAdapter(),
                Duration.ofMillis(20));

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

        Assert.assertEquals(ModelVerificationStatus.PARTIAL, result.getStatus());
        Assert.assertEquals("连接成功，流式响应不可用，可关闭智能体的模型流式响应。",
                result.getMessage());
    }

    /**
     * 验证 VLM 使用 Base64 图片，并能聚合流式文本增量。
     */
    @Test
    public void shouldVerifyVlmImageAndAggregateStreamingChunks() {
        RecordingModelFactory factory = new RecordingModelFactory(
                Flux.just(response("图片中的内容是“" + VlmVerificationImage.VERIFICATION_CODE + "”。")),
                Flux.just(response("识别结果：58"), response("39")));

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

        Assert.assertEquals(ModelVerificationStatus.PASSED, result.getStatus());
        Msg message = factory.getMessages().get(0);
        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]);
    }

    /**
     * 创建待验证模型。
     *
     * @param supportImage 是否支持图片
     * @return 测试模型
     */
    private Model model(boolean supportImage) {
        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);
        ModelProvider provider = new ModelProvider();
        provider.setProviderType("gpustack");
        model.setModelProvider(provider);
        return model;
    }

    /**
     * 创建使用测试工厂的验证器。
     *
     * @param factory 记录型模型工厂
     * @return 验证器
     */
    private AgentScopeChatModelConnectivityVerifier verifier(RecordingModelFactory factory) {
        return new AgentScopeChatModelConnectivityVerifier(
                factory,
                new AgentScopeMessageAdapter(),
                Duration.ofSeconds(2));
    }

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

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

        /** 每个阶段的预设响应。 */
        private final List<Flux<ChatResponse>> phaseResponses;
        /** 模型工厂收到的流式参数。 */
        private final List<Boolean> factoryStreams = new ArrayList<>();
        /** 模型请求收到的流式参数。 */
        private final List<Boolean> requestStreams = new ArrayList<>();
        /** 模型请求收到的消息。 */
        private final List<Msg> messages = new ArrayList<>();
        /** OpenAI-compatible 请求中的思考开关。 */
        private final List<Object> enableThinkingValues = new ArrayList<>();
        /** GPUStack 模板参数中的思考开关。 */
        private final List<Object> chatTemplateThinkingValues = new ArrayList<>();

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

        /**
         * 创建当前验证阶段的测试模型。
         *
         * @param modelSpec 模型声明
         * @param generationOptions 生成参数
         * @return 测试模型
         */
        @Override
        public io.agentscope.core.model.Model create(
                AgentModelSpec modelSpec,
                AgentGenerationOptions generationOptions) {
            int phaseIndex = factoryStreams.size();
            factoryStreams.add(Boolean.TRUE.equals(generationOptions.getStream()));
            enableThinkingValues.add(
                    generationOptions.getAdditionalBodyParams().get("enable_thinking"));
            Object templateOptions = generationOptions.getAdditionalBodyParams()
                    .get("chat_template_kwargs");
            chatTemplateThinkingValues.add(templateOptions instanceof java.util.Map<?, ?> map
                    ? map.get("enable_thinking")
                    : null);
            Flux<ChatResponse> responses = phaseResponses.get(phaseIndex);
            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()));
                    messages.add(inputMessages.get(0));
                    return responses;
                }

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

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

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

        /**
         * 获取请求消息记录。
         *
         * @return 消息列表
         */
        private List<Msg> getMessages() {
            return messages;
        }

        /**
         * 获取 OpenAI-compatible 请求中的思考开关。
         *
         * @return 各阶段思考开关
         */
        private List<Object> getEnableThinkingValues() {
            return enableThinkingValues;
        }

        /**
         * 获取 GPUStack 模板参数中的思考开关。
         *
         * @return 各阶段模板思考开关
         */
        private List<Object> getChatTemplateThinkingValues() {
            return chatTemplateThinkingValues;
        }
    }
}
