package com.easyagents.flow.support.provider;

import com.easyagents.core.message.AiMessage;
import com.easyagents.core.message.SystemMessage;
import com.easyagents.core.model.chat.BaseChatModel;
import com.easyagents.core.model.chat.ChatModel;
import com.easyagents.core.model.chat.StreamResponseListener;
import com.easyagents.core.model.client.StreamContext;
import com.easyagents.core.model.chat.response.AiMessageResponse;
import com.easyagents.core.prompt.SimplePrompt;
import com.easyagents.core.util.ImageUtil;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainStatus;
import com.easyagents.flow.core.chain.event.ChainStatusChangeEvent;
import com.easyagents.flow.core.chain.event.LlmStreamEvent;
import com.easyagents.flow.core.chain.listener.ChainEventListener;
import com.easyagents.flow.core.llm.Llm;
import com.easyagents.flow.core.node.LlmNode;
import com.easyagents.flow.core.util.StringUtil;

import java.io.File;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 基于 Easy-Agents 聊天模型实现工作流 LLM 调用。
 */
public class EasyAgentsLlm implements Llm {

    private ChatModel chatModel;
    private ImageInputResolver imageInputResolver;

    /**
     * 获取聊天模型。
     *
     * @return 聊天模型
     */
    public ChatModel getChatModel() {
        return chatModel;
    }

    /**
     * 设置聊天模型。
     *
     * @param chatModel 聊天模型
     */
    public void setChatModel(ChatModel chatModel) {
        this.chatModel = chatModel;
    }

    /**
     * 获取图片输入解析器。
     *
     * @return 图片输入解析器
     */
    public ImageInputResolver getImageInputResolver() {
        return imageInputResolver;
    }

    /**
     * 设置图片输入解析器。
     *
     * @param imageInputResolver 图片输入解析器
     */
    public void setImageInputResolver(ImageInputResolver imageInputResolver) {
        this.imageInputResolver = imageInputResolver;
    }

    /**
     * 调用聊天模型并返回文本结果。
     *
     * @param messageInfo 消息信息
     * @param options     模型参数
     * @param llmNode     当前 LLM 节点
     * @param chain       工作流链
     * @return 模型文本结果
     */
    @Override
    public String chat(MessageInfo messageInfo, ChatOptions options, LlmNode llmNode, Chain chain) {

        SimplePrompt prompt = buildPrompt(messageInfo);
        com.easyagents.core.model.chat.ChatOptions chatOptions = buildChatOptions(options);
        CountDownLatch completion = new CountDownLatch(1);
        AtomicReference<StreamContext> streamContext = new AtomicReference<>();
        AtomicReference<String> result = new AtomicReference<>();
        AtomicReference<Throwable> failure = new AtomicReference<>();
        String streamId = UUID.randomUUID().toString();

        ChainEventListener cancellationListener = (event, eventChain) -> {
            if (!(event instanceof ChainStatusChangeEvent statusEvent)
                    || statusEvent.getStatus() != ChainStatus.CANCELLED
                    || !chain.getStateInstanceId().equals(eventChain.getStateInstanceId())) {
                return;
            }
            StreamContext context = streamContext.get();
            if (context != null) {
                context.getClient().stop();
            }
        };
        chain.getEventManager().addEventListener(
                ChainStatusChangeEvent.class, cancellationListener);

        try {
            chatModel.chatStream(prompt, new StreamResponseListener() {
                /**
                 * 记录流客户端，供工作流取消时立即关闭模型连接。
                 *
                 * @param context 流上下文
                 */
                @Override
                public void onStart(StreamContext context) {
                    streamContext.set(context);
                }

                /**
                 * 将模型文本增量转发为工作流事件。
                 *
                 * @param context 流上下文
                 * @param response 本次模型响应
                 */
                @Override
                public void onMessage(StreamContext context, AiMessageResponse response) {
                    AiMessage message = response == null ? null : response.getMessage();
                    if (message == null) {
                        return;
                    }
                    String reasoningDelta = message.getReasoningContent();
                    if (StringUtil.hasText(reasoningDelta)) {
                        chain.notifyEvent(new LlmStreamEvent(
                                chain,
                                llmNode,
                                streamId,
                                reasoningDelta,
                                LlmStreamEvent.ContentType.REASONING
                        ));
                    }
                    String textDelta = message.getContent();
                    if (StringUtil.hasText(textDelta)) {
                        chain.notifyEvent(new LlmStreamEvent(
                                chain,
                                llmNode,
                                streamId,
                                textDelta,
                                LlmStreamEvent.ContentType.TEXT
                        ));
                    }
                }

                /**
                 * 收集完整响应并结束同步节点等待。
                 *
                 * @param context 流上下文
                 */
                @Override
                public void onStop(StreamContext context) {
                    try {
                        if (failure.get() == null) {
                            AiMessage message = context.getFullMessage();
                            if (message == null || StringUtil.noText(message.getFullContent())) {
                                failure.compareAndSet(
                                        null,
                                        new IllegalStateException(
                                                "EasyAgentsLlm can not get aiMessage!"));
                            } else {
                                result.set(message.getFullContent());
                            }
                        }
                    } finally {
                        completion.countDown();
                    }
                }

                /**
                 * 记录模型流异常并结束等待。
                 *
                 * @param context 流上下文
                 * @param throwable 模型异常
                 */
                @Override
                public void onFailure(StreamContext context, Throwable throwable) {
                    failure.compareAndSet(
                            null,
                            throwable == null
                                    ? new IllegalStateException("EasyAgentsLlm stream failed")
                                    : throwable);
                    completion.countDown();
                }
            }, chatOptions);
            awaitCompletion(completion, streamContext);
        } finally {
            chain.getEventManager().removeEventListener(
                    ChainStatusChangeEvent.class, cancellationListener);
        }

        Throwable throwable = failure.get();
        if (throwable != null) {
            throw new RuntimeException("EasyAgentsLlm stream failed", throwable);
        }
        if (StringUtil.noText(result.get())) {
            throw new RuntimeException("EasyAgentsLlm can not get response!");
        }
        return result.get();
    }

    /**
     * 构建模型提示词，并解析图片输入。
     *
     * @param messageInfo 消息信息
     * @return 模型提示词
     */
    private SimplePrompt buildPrompt(MessageInfo messageInfo) {
        SimplePrompt prompt = new SimplePrompt(messageInfo.getMessage());

        if (StringUtil.hasText(messageInfo.getSystemMessage())) {
            prompt.setSystemMessage(SystemMessage.of(messageInfo.getSystemMessage()));
        }

        List<String> images = resolveImages(messageInfo);
        if (images != null && !images.isEmpty()) {
            assertImageSupported();
            for (String image : images) {
                prompt.addImageUrl(image);
            }
        }
        return prompt;
    }

    /**
     * 构建 Easy-Agents 模型参数。
     *
     * @param options 工作流模型参数
     * @return Easy-Agents 模型参数
     */
    private com.easyagents.core.model.chat.ChatOptions buildChatOptions(ChatOptions options) {
        com.easyagents.core.model.chat.ChatOptions chatOptions = new com.easyagents.core.model.chat.ChatOptions();
        chatOptions.setSeed(options.getSeed());
        chatOptions.setTemperature(options.getTemperature());
        chatOptions.setTopP(options.getTopP());
        chatOptions.setTopK(options.getTopK());
        chatOptions.setMaxTokens(options.getMaxTokens());
        chatOptions.setStop(options.getStop());
        return chatOptions;
    }

    /**
     * 等待异步模型流结束。
     *
     * @param completion 完成信号
     * @param streamContext 当前模型流上下文
     */
    private void awaitCompletion(
            CountDownLatch completion,
            AtomicReference<StreamContext> streamContext) {
        try {
            completion.await();
        } catch (InterruptedException exception) {
            Thread.currentThread().interrupt();
            StreamContext context = streamContext.get();
            if (context != null) {
                context.getClient().stop();
            }
            throw new RuntimeException("EasyAgentsLlm stream interrupted", exception);
        }
    }

    /**
     * 将原始图片输入解析为模型图片 URL。
     *
     * @param messageInfo 消息信息
     * @return 已解析图片 URL 列表
     */
    List<String> resolveImages(MessageInfo messageInfo) {
        if (messageInfo == null) {
            return Collections.emptyList();
        }
        List<?> inputs = messageInfo.getImageInputs();
        if (inputs == null || inputs.isEmpty()) {
            inputs = messageInfo.getImages();
        }
        if (inputs == null || inputs.isEmpty()) {
            return Collections.emptyList();
        }

        List<String> resolvedImages = new ArrayList<>(inputs.size());
        for (Object input : inputs) {
            if (input == null) {
                continue;
            }
            String resolvedImage = resolveImage(input);
            if (StringUtil.noText(resolvedImage)) {
                throw new IllegalArgumentException("Resolved image input must not be blank");
            }
            resolvedImages.add(resolvedImage);
        }
        return resolvedImages;
    }

    /**
     * 解析单个图片输入。
     *
     * @param imageInput 原始图片输入
     * @return 图片 URL 或 Data URI
     */
    private String resolveImage(Object imageInput) {
        if (imageInputResolver != null) {
            return imageInputResolver.resolve(imageInput);
        }
        if (imageInput instanceof String value) {
            return value;
        }
        if (imageInput instanceof File file) {
            return ImageUtil.imageFileToDataUri(file);
        }
        throw new IllegalArgumentException(
                "Unsupported image input type: " + imageInput.getClass().getName());
    }

    /**
     * 校验当前聊天模型是否明确支持图片。
     */
    private void assertImageSupported() {
        if (chatModel instanceof BaseChatModel<?> baseChatModel
                && Boolean.FALSE.equals(baseChatModel.getConfig().getSupportImage())) {
            throw new IllegalArgumentException("当前模型不支持图片输入，请选择支持视觉能力的模型");
        }
    }
}
