package com.tianji.aigc.agent;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.tianji.aigc.attachment.AttachmentContext;
import com.tianji.aigc.attachment.AttachmentContextHolder;
import com.tianji.aigc.config.ModelOptionsHolder;
import com.tianji.aigc.config.ModelOptionsHolder.ModelOptions;
import com.tianji.aigc.config.ToolResultHolder;
import com.tianji.aigc.constants.Constant;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.harness.HarnessEventRecorder;
import com.tianji.aigc.service.ChatService;
import com.tianji.aigc.service.ChatSessionService;
import com.tianji.aigc.vo.ChatEventVO;
import com.tianji.common.utils.UserContext;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.openai.OpenAiChatOptions;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
public abstract class AbstractAgent implements Agent {

    @Resource
    private ChatClient dashScopeChatClient;
    @Resource
    private ChatClient openAiChatClient;
    @Resource
    private ChatSessionService chatSessionService;
    @Resource
    private ChatMemory chatMemory;
    @Resource(name = "messageChatMemoryAdvisor")
    private Advisor messageChatMemoryAdvisor;

    // 输出结束的标记
    public static final ChatEventVO STOP_EVENT = ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build();

    // 存储大模型的生成状态，这里采用ConcurrentHashMap是确保线程安全
    // 目前的版本暂时用Map实现，如果考虑分布式环境的话，可以考虑用redis来实现
    public static final Map<String, Boolean> GENERATE_STATUS = new ConcurrentHashMap<>();

    @Override
    public String process(String question, String sessionId) {
        // 获取用户id
        var userId = UserContext.getUser();
        var requestId = this.generateRequestId();
        this.registerAttachmentParams(sessionId, requestId);

        //更新会话时间
        this.chatSessionService.update(sessionId, question, userId);
        try {
            return this.getChatClientRequest(sessionId, requestId, question)
                    .call()
                    .content();
        } finally {
            if (this.useAttachmentContext()) {
                AttachmentContextHolder.clear(sessionId);
            }
        }
    }

    public Flux<ChatEventVO> processStream(String question, String sessionId) {
        // 获取用户id
        var userId = UserContext.getUser();
        var requestId = this.generateRequestId();
        // 大模型输出内容的缓存器，用于在输出中断后的数据存储
        StringBuilder outputBuilder = new StringBuilder();
        this.registerAttachmentParams(sessionId, requestId);

        //更新会话时间
        this.chatSessionService.update(sessionId, question, userId);

        return this.getChatClientRequest(sessionId, requestId, question)
                .stream()
                .chatResponse()
                .doFirst(() -> {
                    //输出开始，标记正在输出
                    GENERATE_STATUS.put(sessionId, true);
                })
                .doOnComplete(() -> {
                    //输出结束，清除标记
                    GENERATE_STATUS.remove(sessionId);
                })
                .doOnError(throwable -> {
                    GENERATE_STATUS.remove(sessionId);
                    if (this.useAttachmentContext()) {
                        AttachmentContextHolder.clear(sessionId);
                    }
                }) // 错误时清除标记
                .doOnCancel(() -> {
                    // 当输出被取消时，保存输出的内容到历史记录中
                    this.saveStopHistoryRecord(sessionId, outputBuilder.toString());
                    if (this.useAttachmentContext()) {
                        AttachmentContextHolder.clear(sessionId);
                    }
                })
                .takeWhile(s -> Optional.ofNullable(GENERATE_STATUS.get(sessionId)).orElse(false)) // 只输出标记为true的流
                .map(chatResponse -> {
                    // 对于响应结果进行处理，如果是最后一条数据，就把此次消息id放到内存中
                    // 主要用于存储消息数据到 redis中，可以根据消息di获取的请求id，再通过请求id就可以获取到参数列表了
                    // 从而解决，在历史聊天记录中没有外参数的问题
                    var finishReason = chatResponse.getResult().getMetadata().getFinishReason();
                    if (StrUtil.equals(Constant.STOP, finishReason)) {
                        var messageId = chatResponse.getMetadata().getId();
                        ToolResultHolder.put(messageId, Constant.REQUEST_ID, requestId);
                    }
                    // 获取大模型的输出的内容
                    String text = chatResponse.getResult().getOutput().getText();
                    // 追加到输出内容中
                    outputBuilder.append(text);
                    // 封装响应对象
                    return ChatEventVO.builder()
                            .eventData(text)
                            .eventType(ChatEventTypeEnum.DATA.getValue())
                            .build();
                })
                .concatWith(Flux.defer(() -> {
                    // 通过请求id获取到参数列表，如果不为空，就将其追加到返回结果中
                    try {
                        var map = ToolResultHolder.get(requestId);
                        if (CollUtil.isNotEmpty(map)) {
                            ToolResultHolder.remove(requestId); // 清除参数列表
                            List<ChatEventVO> events = new ArrayList<>();
                            Object trace = map.get(HarnessEventRecorder.TRACE_FIELD);
                            if (trace != null) {
                                events.add(ChatEventVO.builder()
                                        .eventData(trace)
                                        .eventType(ChatEventTypeEnum.TRACE.getValue())
                                        .build());
                            }
                            // 响应给前端的参数数据
                            ChatEventVO chatEventVO = ChatEventVO.builder()
                                    .eventData(map)
                                    .eventType(ChatEventTypeEnum.PARAM.getValue())
                                    .build();
                            events.add(chatEventVO);
                            events.add(STOP_EVENT);
                            return Flux.fromIterable(events);
                        }
                        return Flux.just(STOP_EVENT);
                    }
                    finally {
                        if (this.useAttachmentContext()) {
                            AttachmentContextHolder.clear(sessionId);
                        }
                    }
                }));
    }

    /**
     * 保存停止输出的记录
     *
     * @param sessionId 会话id
     * @param content   大模型输出的内容
     */
    private void saveStopHistoryRecord(String sessionId, String content) {
        String conversationId = ChatService.getConversationId(sessionId);
        this.chatMemory.add(conversationId, new AssistantMessage(content));
    }

    private String generateRequestId() {
        return IdUtil.fastSimpleUUID();
    }

    private ChatClient.ChatClientRequestSpec getChatClientRequest(String sessionId, String requestId, String question) {
        List<Advisor> advisors = new ArrayList<>();
        if (this.useChatMemory()) {
            advisors.add(this.messageChatMemoryAdvisor);
        }
        advisors.addAll(this.advisors(question));

        ChatClient client = resolveChatClient();
        ChatClient.ChatClientRequestSpec request = client.prompt()
                .system(promptSystem -> promptSystem.text(this.resolveSystemMessage(sessionId)).params(this.systemMessageParams()))
                .advisors(advisor -> advisor.advisors(advisors).params(this.advisorParams(sessionId, requestId)))
                .tools(this.tools())
                .toolContext(this.toolContext(sessionId, requestId))
                .user(question);

        applyChatOptions(request);
        return request;
    }

    /**
     * 根据 ModelOptionsHolder 中的 provider 选择对应的 ChatClient。
     * 默认使用 dashScopeChatClient。
     */
    private ChatClient resolveChatClient() {
        ModelOptions options = ModelOptionsHolder.get();
        if (options != null && options.hasProvider() && "openai".equalsIgnoreCase(options.provider())) {
            return this.openAiChatClient;
        }
        return this.dashScopeChatClient;
    }

    /**
     * 从 ModelOptionsHolder 读取 model 和 temperature，构建对应的 ChatOptions 并应用到请求。
     */
    private void applyChatOptions(ChatClient.ChatClientRequestSpec request) {
        ModelOptions options = ModelOptionsHolder.get();
        if (options == null) return;

        String provider = options.hasProvider() ? options.provider() : "dashscope";

        if ("openai".equalsIgnoreCase(provider)) {
            OpenAiChatOptions oaiOptions = new OpenAiChatOptions();
            if (options.hasModel()) oaiOptions.setModel(options.model());
            if (options.hasTemperature()) oaiOptions.setTemperature(options.temperature());
            request.options(oaiOptions);
        } else {
            DashScopeChatOptions dsOptions = new DashScopeChatOptions();
            if (options.hasModel()) dsOptions.setModel(options.model());
            if (options.hasTemperature()) dsOptions.setTemperature(options.temperature());
            request.options(dsOptions);
        }
    }

    protected boolean useAttachmentContext() {
        return true;
    }

    private void registerAttachmentParams(String sessionId, String requestId) {
        if (!this.useAttachmentContext()) {
            return;
        }
        AttachmentContext context = AttachmentContextHolder.peek(sessionId);
        if (context != null && context.hasSources()) {
            ToolResultHolder.putAll(requestId, context.toParamMap());
        }
    }

    private String resolveSystemMessage(String sessionId) {
        String systemMessage = this.systemMessage();
        if (!this.useAttachmentContext()) {
            return systemMessage;
        }
        AttachmentContext context = AttachmentContextHolder.peek(sessionId);
        if (context == null || !context.hasSources()) {
            return systemMessage;
        }
        return systemMessage + "\n\n" + context.toSystemPrompt();
    }

    @Override
    public Map<String, Object> advisorParams(String sessionId, String requestId) {
        String conversationId = ChatService.getConversationId(sessionId);
        return Map.of(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId);
    }

    @Override
    public void stop(String sessionId) {
        GENERATE_STATUS.remove(sessionId);
    }
}
