package cc.mrbird.febs.ai.strategy.Impl;

import cc.mrbird.febs.ai.req.agent.AiRequestDto;
import cc.mrbird.febs.ai.strategy.LlmStrategyService;
import cc.mrbird.febs.ai.strategy.enumerates.LlmStrategyContextEnum;
import cc.mrbird.febs.ai.strategy.param.LlmStrategyDto;
import cc.mrbird.febs.common.entity.FebsResponse;
import cc.mrbird.febs.common.exception.FebsException;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.app.*;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.JsonUtils;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import javax.annotation.PostConstruct;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

@Slf4j
@Component("AliLlmStrategyService")
public class AliLlmStrategyServiceImpl implements LlmStrategyService {

    private GenerationParam generationParam;

    private static final String apiKey = "sk-a2323eba1e584066b3a536aefa804970";
    private static final String model = "qwen-plus";


    private ApplicationParam applicationParam;
    private ApplicationParam applicationParamSummary;

    private static final String appId = "a5d38240f6b94da6b7f1f7fbe563f50c";
    private static final String SummaryAppId = "19da84392c534db5b2e3f10e698758ab";

    @PostConstruct
    public void init() {
        this.generationParam = GenerationParam.builder()
                // 若没有配置环境变量，请用阿里云百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey(apiKey)
                // 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
                .model(model)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .build();

        this.applicationParam = ApplicationParam.builder()
                // 若没有配置环境变量，可用百炼API Key将下行替换为：.apiKey("sk-xxx")。但不建议在生产环境中直接将API Key硬编码到代码中，以减少API Key泄露风险。
                .apiKey(apiKey)
                .appId(appId)
                .build();

        this.applicationParamSummary = ApplicationParam.builder()
                // 若没有配置环境变量，可用百炼API Key将下行替换为：.apiKey("sk-xxx")。但不建议在生产环境中直接将API Key硬编码到代码中，以减少API Key泄露风险。
                .apiKey(apiKey)
                .appId(SummaryAppId)
                .build();
    }

    private List<Message> getMessages(List<LlmStrategyDto> dto) {
        List<Message> messages = new ArrayList<>();
        for (LlmStrategyDto dtoItem : dto){
            if (StrUtil.equals(dtoItem.getRole(),Role.SYSTEM.getValue())){
                messages.add(Message.builder()
                        .role(Role.SYSTEM.getValue())
                        .content(dtoItem.getContent())
                        .build());
            }
            if (StrUtil.equals(dtoItem.getRole(),Role.USER.getValue())){
                messages.add(Message.builder()
                        .role(Role.USER.getValue())
                        .content(dtoItem.getContent())
                        .build());
            }
            if (StrUtil.equals(dtoItem.getRole(),Role.ASSISTANT.getValue())){
                messages.add(Message.builder()
                        .role(Role.ASSISTANT.getValue())
                        .content(dtoItem.getContent())
                        .build());
            }
        }
        return messages;
    }

    @Override
    public FebsResponse llmInvokeNonStreaming(List<LlmStrategyDto> dto) {
        if (CollUtil.isEmpty(dto)){
            throw new FebsException("百炼大模型初始化异常");
        }
        List<Message> messages = getMessages(dto);
        Generation gen = new Generation();
        generationParam.setMessages(messages);
        generationParam.setEnableThinking( false);
        FebsResponse febsResponse = new FebsResponse();
        try {
            GenerationResult result = gen.call(generationParam);
            if (result != null && result.getOutput() != null && result.getOutput().getChoices().size() > 0){
                febsResponse.success().data(result.getOutput().getChoices().get(0).getMessage().getContent());
            }else{
                febsResponse.fail().message("百炼大模型调用失败");
            }
        } catch (NoApiKeyException e) {
            throw new FebsException(StrUtil.format("百炼大模型调用失败:{}",e.getMessage()));
        } catch (InputRequiredException e) {
            throw new FebsException(StrUtil.format("百炼大模型输出失败:{}",e.getMessage()));
        }
        return febsResponse;
    }

    @Override
    public Flux<FebsResponse> llmInvokeStreamingWithThink(List<LlmStrategyDto> dto) {
        if (CollUtil.isEmpty(dto)){
            throw new FebsException("百炼大模型初始化异常");
        }
        List<Message> messages = getMessages(dto);

        long startTime = System.currentTimeMillis();
        Generation gen = new Generation();
        generationParam.setMessages(messages);
        generationParam.setResultFormat(GenerationParam.ResultFormat.MESSAGE);
        generationParam.setEnableThinking( true);
        generationParam.setIncrementalOutput( true);
        Flowable<GenerationResult> result;
        try {
            result = gen.streamCall(generationParam);
        } catch (NoApiKeyException | InputRequiredException e) {
            throw new FebsException(StrUtil.format("百炼大模型输出失败:{}",e.getMessage()));
        }

        return Flux.from(result)
                .map(message -> {
                    HashMap<String, String> stringStringHashMap = new HashMap<>();
                    if (StrUtil.isNotEmpty(message.getOutput().getChoices().get(0).getMessage().getReasoningContent())){
                        stringStringHashMap.put(LlmStrategyContextEnum.THINK.name(),message.getOutput().getChoices().get(0).getMessage().getReasoningContent());

                        System.out.print(message.getOutput().getChoices().get(0).getMessage().getReasoningContent());
                    }
                    if (StrUtil.isNotEmpty(message.getOutput().getChoices().get(0).getMessage().getContent())){
                        stringStringHashMap.put(LlmStrategyContextEnum.CONTENT.name(),message.getOutput().getChoices().get(0).getMessage().getContent());
                        System.out.print(message.getOutput().getChoices().get(0).getMessage().getContent());
                    }
                    return new FebsResponse().success().data(stringStringHashMap);
                })
                .doOnComplete(() -> {
                    long endTime = System.currentTimeMillis();
                    System.out.println("百炼大模型输出：" + (endTime - startTime) + "毫秒");
                })
                .doOnError(error -> {
                    throw new FebsException(StrUtil.format("百炼大模型输出失败:{}",error));
                });
    }

    @Override
    public Flux<FebsResponse> llmInvokeStreamingNoThink(List<LlmStrategyDto> dto) {
        if (CollUtil.isEmpty(dto)){
            throw new FebsException("百炼大模型初始化异常");
        }
        List<Message> messages = getMessages(dto);

        long startTime = System.currentTimeMillis();
        Generation gen = new Generation();
        generationParam.setMessages(messages);
        generationParam.setResultFormat(GenerationParam.ResultFormat.MESSAGE);
        generationParam.setIncrementalOutput( true);
        generationParam.setEnableThinking( false);
        Flowable<GenerationResult> result;
        try {
            result = gen.streamCall(generationParam);
        } catch (NoApiKeyException | InputRequiredException e) {
            throw new FebsException(StrUtil.format("百炼大模型输出失败:{}",e.getMessage()));
        }

        return Flux.from(result)
                .map(message -> {
                    HashMap<String, String> stringStringHashMap = new HashMap<>();
                    if (StrUtil.isNotEmpty(message.getOutput().getChoices().get(0).getMessage().getContent())){
                        String content = message.getOutput().getChoices().get(0).getMessage().getContent();
                        System.out.print( content);
                        stringStringHashMap.put(LlmStrategyContextEnum.CONTENT.name(),content);
                    }
                    return new FebsResponse().success().data(stringStringHashMap);
                })
                .doOnComplete(() -> {
                    long endTime = System.currentTimeMillis();
                    System.out.println("百炼大模型输出：" + (endTime - startTime) + "毫秒");
                })
                .doOnError(error -> {
                    throw new FebsException(StrUtil.format("百炼大模型输出失败:{}",error));
                });
    }

    @Override
    public Flux<FebsResponse> llmInvokeStreamingNoThink(AiRequestDto dto) {
        long startTime = System.currentTimeMillis();
        if (ObjectUtil.isEmpty(dto)){
            throw new FebsException("参数异常");
        }
        String talkId = dto.getTalkId();

        String rolePrompt = dto.getRolePrompt();
        List<Message> messages = dto.getMessages();
        HashMap<String, Object> userPromptParams = new HashMap<>();
        userPromptParams.put("role_prompt", rolePrompt);
        if(CollUtil.isNotEmpty( messages)){
            userPromptParams.put("message_list", messages);
        }
        HashMap<String, Object> bizParams = new HashMap<>();
        bizParams.put("user_prompt_params", userPromptParams);

        List<String> knowledgeIds = dto.getKnowledgeIds();
        List<String> fileIds = dto.getFileIds();
        String prompt = dto.getPrompt();

        Application application = new Application();
        applicationParam.setPrompt(prompt);
        applicationParam.setBizParams(JsonUtils.toJsonObject( bizParams));
        // 获取当前的系统时间
        applicationParam.setEnableSystemTime( true);
        // 增量输出
        applicationParam.setIncrementalOutput( true);
        // 思考过程
        applicationParam.setEnableThinking( false);
        // 多轮对话
//        applicationParam.setMessages(messages);
        // 上下文seesionId
//        String sessionId = (String) redisCache.getCacheObject(talkId);
//        if (ObjectUtil.isNotEmpty(sessionId)){
//            applicationParam.setSessionId(talkId);
//        }

        RagOptions ragOptions = null;
        if (CollUtil.isEmpty(fileIds)){
            ragOptions = RagOptions.builder()
                    .pipelineIds(knowledgeIds)
                    .build();
        }else{
            ragOptions = RagOptions.builder()
                    .pipelineIds(knowledgeIds)
                    .fileIds(fileIds)
                    .build();
        }
        applicationParam.setRagOptions(ragOptions);
        Flowable<ApplicationResult> result;
        try {
            result = application.streamCall(applicationParam);
        } catch (NoApiKeyException | InputRequiredException e) {

            throw new FebsException(StrUtil.format("百炼大模型输出失败:{}",e.getMessage()));
        }
        return Flux.from(result)
                .map(message -> {
                    HashMap<String, Object> stringStringHashMap = new HashMap<>();
                    if (!message.getOutput().getFinishReason().equals("stop")){
                        stringStringHashMap.put(LlmStrategyContextEnum.CONTENT.name(),message.getOutput().getText());
                    }
                    if (message.getOutput().getFinishReason().equals("stop")){
                        log.info("百炼大模型输出：{}",message.getOutput().getSessionId());
                        List<ApplicationUsage.ModelUsage> models = message.getUsage().getModels();
                        long outputTokens = models.stream().mapToLong(ApplicationUsage.ModelUsage::getOutputTokens).sum();
                        log.info("百炼大模型输出：{}",outputTokens);
                        long inputTokens = models.stream().mapToLong(ApplicationUsage.ModelUsage::getInputTokens).sum();
                        log.info("百炼大模型输出：{}",inputTokens);
                        stringStringHashMap.put("inputTokens:",inputTokens);
                        stringStringHashMap.put("outputTokens:",outputTokens);
                        stringStringHashMap.put(LlmStrategyContextEnum.CONTENT.name(),message.getOutput().getText());
                    }

                    return new FebsResponse().success().data(stringStringHashMap);
                })
                .doOnComplete(() -> {
                    long endTime = System.currentTimeMillis();
                    System.out.println("百炼大模型输出：" + (endTime - startTime) + "毫秒");
                })
                .doOnError(error -> {
                    throw new FebsException(StrUtil.format("百炼大模型输出失败:{}",error));
                });
    }
}