package com.easyagents.flow.support.provider;

import com.easyagents.core.model.chat.BaseChatModel;
import com.easyagents.core.model.chat.ChatConfig;
import com.easyagents.core.model.chat.ChatModel;
import com.easyagents.core.model.chat.ChatOptions;
import com.easyagents.core.model.chat.StreamResponseListener;
import com.easyagents.core.model.chat.response.AiMessageResponse;
import com.easyagents.core.model.client.StreamClient;
import com.easyagents.core.model.client.StreamContext;
import com.easyagents.core.prompt.Prompt;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainDefinition;
import com.easyagents.flow.core.chain.EventManager;
import com.easyagents.flow.core.chain.event.LlmStreamEvent;
import com.easyagents.flow.core.llm.Llm;
import com.easyagents.flow.core.node.LlmNode;
import org.junit.Assert;
import org.junit.Test;

import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * 工作流 LLM 图片输入解析测试。
 */
public class EasyAgentsLlmTest {

    /**
     * 验证结构化图片对象通过注入的解析器转换。
     */
    @Test
    public void shouldResolveRawImageInput() {
        EasyAgentsLlm llm = new EasyAgentsLlm();
        llm.setImageInputResolver(input -> {
            Assert.assertTrue(input instanceof Map<?, ?>);
            return "data:image/png;base64,AQID";
        });
        Llm.MessageInfo messageInfo = new Llm.MessageInfo();
        messageInfo.setImageInputs(List.of(Map.of(
                "sourceType", "upload",
                "fileName", "image.png",
                "filePath", "workflow/image.png")));

        Assert.assertEquals(
                List.of("data:image/png;base64,AQID"),
                llm.resolveImages(messageInfo));
    }

    /**
     * 验证旧版图片字符串仍然可直接传递。
     */
    @Test
    public void shouldKeepLegacyImageString() {
        EasyAgentsLlm llm = new EasyAgentsLlm();
        Llm.MessageInfo messageInfo = new Llm.MessageInfo();
        messageInfo.setImages(List.of("https://example.com/image.png"));

        Assert.assertEquals(
                List.of("https://example.com/image.png"),
                llm.resolveImages(messageInfo));
    }

    /**
     * 验证 File 图片会补全 MIME 和 Data URI 前缀。
     *
     * @throws Exception 临时文件写入失败
     */
    @Test
    public void shouldConvertFileToCompleteDataUri() throws Exception {
        Path image = Files.createTempFile("easy-agents-image-", ".png");
        try {
            Files.write(image, new byte[]{1, 2, 3});
            EasyAgentsLlm llm = new EasyAgentsLlm();
            Llm.MessageInfo messageInfo = new Llm.MessageInfo();
            messageInfo.setImageInputs(List.of(image.toFile()));

            Assert.assertEquals(
                    List.of("data:image/png;base64,AQID"),
                    llm.resolveImages(messageInfo));
        } finally {
            Files.deleteIfExists(image);
        }
    }

    /**
     * 验证明确不支持图片的模型会在发送请求前失败。
     */
    @Test
    public void shouldRejectImageForUnsupportedModel() {
        ChatConfig config = new ChatConfig();
        config.setSupportImage(Boolean.FALSE);
        EasyAgentsLlm llm = new EasyAgentsLlm();
        llm.setChatModel(new BaseChatModel<>(config) {
        });
        Llm.MessageInfo messageInfo = new Llm.MessageInfo();
        messageInfo.setImages(List.of("data:image/png;base64,AQID"));

        try {
            llm.chat(messageInfo, new Llm.ChatOptions(), null, null);
            Assert.fail("expected IllegalArgumentException");
        } catch (IllegalArgumentException exception) {
            Assert.assertEquals(
                    "当前模型不支持图片输入，请选择支持视觉能力的模型",
                    exception.getMessage());
        }
    }

    /**
     * 验证重复执行同一 LLM 节点时，每次调用拥有独立流标识且增量不会被覆盖。
     */
    @Test
    public void shouldEmitIndependentStreamForRepeatedNodeInvocation() {
        EasyAgentsLlm llm = new EasyAgentsLlm();
        llm.setChatModel(new DeterministicStreamChatModel());

        Chain chain = new Chain(new ChainDefinition(), "stream-test");
        EventManager eventManager = new EventManager();
        chain.setEventManager(eventManager);
        List<LlmStreamEvent> events = new ArrayList<>();
        eventManager.addEventListener(
                LlmStreamEvent.class,
                (event, currentChain) -> events.add((LlmStreamEvent) event));

        LlmNode node = new LlmNode();
        node.setId("llm-1");
        Llm.MessageInfo messageInfo = new Llm.MessageInfo();
        messageInfo.setMessage("请回答");

        Assert.assertEquals(
                "你好",
                llm.chat(messageInfo, new Llm.ChatOptions(), node, chain));
        Assert.assertEquals(
                "你好",
                llm.chat(messageInfo, new Llm.ChatOptions(), node, chain));

        Assert.assertEquals(4, events.size());
        Assert.assertEquals(List.of("你", "好", "你", "好"), events.stream()
                .map(LlmStreamEvent::getDelta)
                .toList());
        Assert.assertTrue(events.stream().noneMatch(LlmStreamEvent::isReasoning));
        Assert.assertEquals(events.get(0).getStreamId(), events.get(1).getStreamId());
        Assert.assertEquals(events.get(2).getStreamId(), events.get(3).getStreamId());
        Assert.assertNotEquals(events.get(0).getStreamId(), events.get(2).getStreamId());
    }

    /**
     * 验证模型思考与正式回答使用同一流标识并按内容类型分别发布。
     */
    @Test
    public void shouldEmitReasoningAndTextDeltasSeparately() {
        EasyAgentsLlm llm = new EasyAgentsLlm();
        llm.setChatModel(new ReasoningStreamChatModel());

        Chain chain = new Chain(new ChainDefinition(), "reasoning-stream-test");
        EventManager eventManager = new EventManager();
        chain.setEventManager(eventManager);
        List<LlmStreamEvent> events = new ArrayList<>();
        eventManager.addEventListener(
                LlmStreamEvent.class,
                (event, currentChain) -> events.add((LlmStreamEvent) event));

        LlmNode node = new LlmNode();
        node.setId("llm-reasoning");
        Llm.MessageInfo messageInfo = new Llm.MessageInfo();
        messageInfo.setMessage("请回答");

        Assert.assertEquals(
                "答案",
                llm.chat(messageInfo, new Llm.ChatOptions(), node, chain));
        Assert.assertEquals(4, events.size());
        Assert.assertEquals(
                List.of("先", "想", "答", "案"),
                events.stream().map(LlmStreamEvent::getDelta).toList());
        Assert.assertEquals(
                List.of(
                        LlmStreamEvent.ContentType.REASONING,
                        LlmStreamEvent.ContentType.REASONING,
                        LlmStreamEvent.ContentType.TEXT,
                        LlmStreamEvent.ContentType.TEXT
                ),
                events.stream().map(LlmStreamEvent::getContentType).toList());
        Assert.assertEquals(
                1,
                events.stream().map(LlmStreamEvent::getStreamId).distinct().count());
    }

    /**
     * 固定输出两个文本增量的测试聊天模型。
     */
    private static final class DeterministicStreamChatModel implements ChatModel {

        /**
         * 同步聊天接口不参与本测试。
         *
         * @param prompt 提示词
         * @param options 模型参数
         * @return 无
         */
        @Override
        public AiMessageResponse chat(Prompt prompt, ChatOptions options) {
            throw new UnsupportedOperationException("sync chat is not used");
        }

        /**
         * 连续发送两个增量及完整消息。
         *
         * @param prompt 提示词
         * @param listener 流监听器
         * @param options 模型参数
         */
        @Override
        public void chatStream(
                Prompt prompt,
                StreamResponseListener listener,
                ChatOptions options) {
            StreamContext context = new StreamContext(this, null, new NoopStreamClient());
            listener.onStart(context);
            for (String content : List.of("你", "好")) {
                com.easyagents.core.message.AiMessage delta =
                        new com.easyagents.core.message.AiMessage();
                delta.setContent(content);
                listener.onMessage(context, new AiMessageResponse(null, content, delta));
            }
            context.setFullMessage(new com.easyagents.core.message.AiMessage("你好"));
            listener.onStop(context);
        }
    }

    /**
     * 固定输出思考与回答增量的测试聊天模型。
     */
    private static final class ReasoningStreamChatModel implements ChatModel {

        /**
         * 同步聊天接口不参与本测试。
         *
         * @param prompt 提示词
         * @param options 模型参数
         * @return 无
         */
        @Override
        public AiMessageResponse chat(Prompt prompt, ChatOptions options) {
            throw new UnsupportedOperationException("sync chat is not used");
        }

        /**
         * 依次发送思考增量、回答增量及完整消息。
         *
         * @param prompt 提示词
         * @param listener 流监听器
         * @param options 模型参数
         */
        @Override
        public void chatStream(
                Prompt prompt,
                StreamResponseListener listener,
                ChatOptions options) {
            StreamContext context = new StreamContext(this, null, new NoopStreamClient());
            listener.onStart(context);
            for (String reasoning : List.of("先", "想")) {
                com.easyagents.core.message.AiMessage delta =
                        new com.easyagents.core.message.AiMessage();
                delta.setReasoningContent(reasoning);
                listener.onMessage(context, new AiMessageResponse(null, reasoning, delta));
            }
            for (String content : List.of("答", "案")) {
                com.easyagents.core.message.AiMessage delta =
                        new com.easyagents.core.message.AiMessage();
                delta.setContent(content);
                listener.onMessage(context, new AiMessageResponse(null, content, delta));
            }
            com.easyagents.core.message.AiMessage fullMessage =
                    new com.easyagents.core.message.AiMessage("答案");
            fullMessage.setFullReasoningContent("先想");
            context.setFullMessage(fullMessage);
            listener.onStop(context);
        }
    }

    /**
     * 测试使用的空流客户端。
     */
    private static final class NoopStreamClient implements StreamClient {

        /**
         * 测试模型直接分发事件，无需启动网络请求。
         *
         * @param url 请求地址
         * @param headers 请求头
         * @param payload 请求体
         * @param listener 客户端监听器
         * @param config 模型配置
         */
        @Override
        public void start(
                String url,
                Map<String, String> headers,
                String payload,
                com.easyagents.core.model.client.StreamClientListener listener,
                ChatConfig config) {
        }

        /**
         * 测试客户端没有需要关闭的网络资源。
         */
        @Override
        public void stop() {
        }
    }
}
