package com.pavis.admin.aigc.core.llm.provider;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.util.ObjectUtil;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.pavis.admin.aigc.constant.ModelConst;
import com.pavis.admin.aigc.core.llm.provider.build.ModelBuildHandler;
import com.pavis.admin.aigc.enums.ModelTypeEnum;
import com.pavis.admin.aigc.mapper.ModelMapper;
import com.pavis.admin.aigc.model.entity.ModelDO;
import com.pavis.admin.aigc.model.query.ModelQuery;
import com.pavis.admin.aigc.model.resp.ModelResp;
import com.pavis.admin.aigc.service.ModelService;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.image.ImageModel;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.annotation.Async;
import top.continew.starter.extension.crud.model.query.SortQuery;

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

/**
 * @author tycoding
 * @since 2024/6/16
 */
@Configuration
@Slf4j
public class ModelStoreFactory {

    @Resource
    private ModelService aigcModelService;

    @Resource
    private ModelMapper modelMapper;
    @Resource
    private List<ModelBuildHandler> modelBuildHandlers;

    private final List<ModelResp> modelStore = new ArrayList<>();
    private final Map<String, StreamingChatLanguageModel> streamingChatMap = new ConcurrentHashMap<>();
    private final Map<String, ChatLanguageModel> chatLanguageMap = new ConcurrentHashMap<>();
    private final Map<String, EmbeddingModel> embeddingModelMap = new ConcurrentHashMap<>();
    private final Map<String, ImageModel> imageModelMap = new ConcurrentHashMap<>();

    @Async
    @PostConstruct
    public void init() {
        modelStore.clear();
        streamingChatMap.clear();
        chatLanguageMap.clear();
        embeddingModelMap.clear();
        imageModelMap.clear();

        List<ModelDO> modelDOS = modelMapper.selectList(new QueryWrapper<>());
        List<ModelResp> list = BeanUtil.copyToList(modelDOS, ModelResp.class);
        // List<ModelResp> list = aigcModelService.list(new ModelQuery(),new SortQuery());
        list.forEach(model -> {
            if (Objects.equals(model.getBaseUrl(), "")) {
                model.setBaseUrl(null);
            }

            chatHandler(model);
            embeddingHandler(model);
            imageHandler(model);
        });

        modelStore.forEach(i -> log.info("已成功注册模型：{} -- {}， 模型配置：{}", i.getProvider(), i.getType(), i));
    }

    private void chatHandler(ModelResp model) {
        try {
            String type = model.getType();
            if (!ModelTypeEnum.CHAT.name().equals(type)) {
                return;
            }
            modelBuildHandlers.forEach(x -> {
                StreamingChatLanguageModel streamingChatLanguageModel = x.buildStreamingChat(model);
                if (ObjectUtil.isNotEmpty(streamingChatLanguageModel)) {
                    streamingChatMap.put(model.getId().toString(), streamingChatLanguageModel);
                    modelStore.add(model);
                }

                ChatLanguageModel languageModel = x.buildChatLanguageModel(model);
                if (ObjectUtil.isNotEmpty(languageModel)) {
                    chatLanguageMap.put(model.getId() + ModelConst.TEXT_SUFFIX, languageModel);
                }
            });
        } catch (Exception e) {
            log.error("model 【 id: {} name: {}】streaming chat 配置报错", model.getId(), model.getName());
        }
    }

    private void embeddingHandler(ModelResp model) {
        try {
            String type = model.getType();
            if (!ModelTypeEnum.EMBEDDING.name().equals(type)) {
                return;
            }
            modelBuildHandlers.forEach(x -> {
                EmbeddingModel embeddingModel = x.buildEmbedding(model);
                if (ObjectUtil.isNotEmpty(embeddingModel)) {
                    embeddingModelMap.put(model.getId().toString(), embeddingModel);
                    modelStore.add(model);
                }
            });

        } catch (Exception e) {
            log.error("model 【id{} name{}】 embedding 配置报错", model.getId(), model.getName());
        }
    }

    private void imageHandler(ModelResp model) {
        try {
            String type = model.getType();
            if (!ModelTypeEnum.TEXT_IMAGE.name().equals(type)) {
                return;
            }
            modelBuildHandlers.forEach(x -> {
                ImageModel imageModel = x.buildImage(model);
                if (ObjectUtil.isNotEmpty(imageModel)) {
                    imageModelMap.put(model.getId().toString(), imageModel);
                    modelStore.add(model);
                }
            });
        } catch (Exception e) {
            log.error("model 【id{} name{}】 image 配置报错", model.getId(), model.getName());
        }
    }

    public StreamingChatLanguageModel getStreamingChatModel(String modelId) {
        return streamingChatMap.get(modelId);
    }

    public boolean containsStreamingChatModel(String modelId) {
        return streamingChatMap.containsKey(modelId);
    }

    public ChatLanguageModel getChatLanguageModel(String modelId) {
        return chatLanguageMap.get(modelId + ModelConst.TEXT_SUFFIX);
    }

    public boolean containsChatLanguageModel(String modelId) {
        return chatLanguageMap.containsKey(modelId + ModelConst.TEXT_SUFFIX);
    }

    public EmbeddingModel getEmbeddingModel(String modelId) {
        return embeddingModelMap.get(modelId);
    }

    public boolean containsEmbeddingModel(String modelId) {
        return embeddingModelMap.containsKey(modelId);
    }

    public ImageModel getImageModel(String modelId) {
        return imageModelMap.get(modelId);
    }

    public boolean containsImageModel(String modelId) {
        return imageModelMap.containsKey(modelId);
    }
}
