package com.pavis.admin.aigc.service.impl;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.pavis.admin.aigc.common.model.req.ChatReq;
import com.pavis.admin.aigc.common.model.req.ImageReq;
import com.pavis.admin.aigc.common.utils.PromptUtil;
import com.pavis.admin.aigc.core.llm.provider.EmbeddingProvider;
import com.pavis.admin.aigc.core.llm.provider.ModelProvider;
import com.pavis.admin.aigc.core.llm.service.Agent;
import com.pavis.admin.aigc.mapper.AppKnowledgeMapper;
import com.pavis.admin.aigc.mapper.KnowledgeMapper;
import com.pavis.admin.aigc.model.entity.KnowledgeDO;

import com.pavis.admin.aigc.model.resp.KnowledgeResp;
import com.pavis.admin.aigc.model.resp.ModelResp;
import com.pavis.admin.aigc.properties.ChatProps;
import com.pavis.admin.aigc.service.KnowledgeService;
import com.pavis.admin.aigc.service.ModelService;
import dev.langchain4j.data.image.Image;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.image.ImageModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.store.embedding.filter.Filter;
import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import top.continew.starter.core.exception.BusinessException;
import top.continew.starter.extension.crud.service.BaseServiceImpl;
import com.pavis.admin.aigc.mapper.AigcAppMapper;
import com.pavis.admin.aigc.model.entity.AigcAppDO;
import com.pavis.admin.aigc.model.query.AigcAppQuery;
import com.pavis.admin.aigc.model.req.AigcAppReq;
import com.pavis.admin.aigc.model.resp.AigcAppDetailResp;
import com.pavis.admin.aigc.model.resp.AigcAppResp;
import com.pavis.admin.aigc.service.AigcAppService;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

/**
 * AIGC应用配置业务实现
 *
 * @author semi
 * @since 2025/05/26 17:28
 */
@Service
@RequiredArgsConstructor
@Slf4j
public class AigcAppServiceImpl extends BaseServiceImpl<AigcAppMapper, AigcAppDO, AigcAppResp, AigcAppDetailResp, AigcAppQuery, AigcAppReq> implements AigcAppService {

    private static final Map<String, AigcAppDO> appMap = new HashMap<>();
    private final ModelService aigcModelService;
    private final KnowledgeService aigcKnowledgeService;
    private final KnowledgeMapper knowledgeMapper;
    private final AppKnowledgeMapper appKnowledgeMapper;
    // longchat中的
    private final ModelProvider provider;
    private final EmbeddingProvider embeddingProvider;
    private final ChatProps chatProps;

    @Override
    public List<AigcAppDO> list(AigcAppDO data) {
        List<AigcAppDO> list = baseMapper.selectList(Wrappers.<AigcAppDO>lambdaQuery()
            .like(StrUtil.isNotBlank(data.getName()), AigcAppDO::getName, data.getName()));
        // copy--》AigcAppDO实体
        List<AigcAppResp> aigcAppResps = BeanUtil.copyToList(list, AigcAppResp.class);

        Map<Long, List<ModelResp>> modelMap = aigcModelService.list(new ModelResp())
            .stream()
            .collect(Collectors.groupingBy(ModelResp::getId));
        Map<Long, List<KnowledgeDO>> knowledgeMap = knowledgeMapper.selectList(new QueryWrapper<>())
            .stream()
            .collect(Collectors.groupingBy(KnowledgeDO::getId));
        aigcAppResps.forEach(i -> {
            List<ModelResp> models = modelMap.get(i.getModelId());
            if (models != null) {
                i.setModel(models.get(0));
            }
            // 通过appid获取knowledgeIds
            List<Long> appKnowledgeIds = appKnowledgeMapper.selectByKnowledgeIds(i.getId());
            if (appKnowledgeIds != null) {
                List<KnowledgeResp> knowledges = new ArrayList<>();
                appKnowledgeIds.forEach(k -> {
                    List<KnowledgeDO> items = knowledgeMap.get(k);
                    KnowledgeResp knowledgeResp = new KnowledgeResp();
                    BeanUtil.copyProperties(items.get(0), knowledgeResp);
                    if (items != null) {
                        knowledges.add(knowledgeResp);
                    }
                });
                i.setKnowledges(knowledges);
            }
        });
        return list;
    }

    @Override
    public AigcAppDO getById(String id) {
        // AigcAppDO app = baseMapper.selectById(id);
        // if (app != null) {
        //     String modelId = app.getModelId();
        //     if (modelId != null) {
        //         app.setModel(aigcModelService.selectById(modelId));
        //     }
        //     List<String> knowledgeIds = app.getKnowledgeIds();
        //     if (knowledgeIds != null && !knowledgeIds.isEmpty()) {
        //         app.setKnowledges(aigcKnowledgeService.list(Wrappers.<AigcKnowledge>lambdaQuery().in(AigcKnowledge::getId, knowledgeIds)));
        //     }
        // }
        // return app;
        return null;
    }

    @PostConstruct
    public void init() {
        log.info("initialize app config list...");
        List<AigcAppDO> list = baseMapper.selectList(new QueryWrapper<>());
        list.forEach(i -> appMap.put(i.getId().toString(), i));
    }

    public AigcAppDO get(String appId) {
        return appMap.get(appId);
    }

    private AiServices<Agent> build(StreamingChatLanguageModel streamModel, ChatLanguageModel model, ChatReq req) {
        AiServices<Agent> aiServices = AiServices.builder(Agent.class)
            .chatMemoryProvider(memoryId -> MessageWindowChatMemory.builder()
                .id(req.getConversationId())
                .chatMemoryStore(new PersistentChatMemoryStore())
                .maxMessages(chatProps.getMemoryMaxMessage())
                .build());
        if (StrUtil.isNotBlank(req.getPromptText())) {
            aiServices.systemMessageProvider(memoryId -> req.getPromptText());
        }
        if (streamModel != null) {
            aiServices.streamingChatLanguageModel(streamModel);
        }
        if (model != null) {
            aiServices.chatLanguageModel(model);
        }
        return aiServices;
    }

    @Override
    public TokenStream chat(ChatReq req) {
        StreamingChatLanguageModel model = provider.stream(req.getModelId());
        if (StrUtil.isBlank(req.getConversationId())) {
            req.setConversationId(IdUtil.simpleUUID());
        }

        AiServices<Agent> aiServices = build(model, null, req);

        if (StrUtil.isNotBlank(req.getKnowledgeId())) {
            req.getKnowledgeIds().add(Long.valueOf(req.getKnowledgeId()));
        }

        if (req.getKnowledgeIds() != null && !req.getKnowledgeIds().isEmpty()) {
            Function<Query, Filter> filter = (query) -> metadataKey("knowledgeId").isIn(req.getKnowledgeIds());
            ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingProvider.getEmbeddingStore(req.getKnowledgeIds()))
                .embeddingModel(embeddingProvider.getEmbeddingModel(req.getKnowledgeIds()))
                .dynamicFilter(filter)
                .build();
            aiServices.retrievalAugmentor(DefaultRetrievalAugmentor.builder()
                .contentRetriever(contentRetriever)
                .build());
        }
        Agent agent = aiServices.build();
        return agent.stream(req.getConversationId(), req.getMessage());
    }

    @Override
    public TokenStream singleChat(ChatReq req) {
        StreamingChatLanguageModel model = provider.stream(req.getModelId());
        if (StrUtil.isBlank(req.getConversationId())) {
            req.setConversationId(IdUtil.simpleUUID());
        }

        Agent agent = build(model, null, req).build();
        if (req.getPrompt() == null) {
            req.setPrompt(PromptUtil.build(req.getMessage(), req.getPromptText()));
        }
        return agent.stream(req.getConversationId(), req.getPrompt().text());
    }

    @Override
    public String text(ChatReq req) {
        if (StrUtil.isBlank(req.getConversationId())) {
            req.setConversationId(IdUtil.simpleUUID());
        }

        try {
            ChatLanguageModel model = provider.text(req.getModelId());
            Agent agent = build(null, model, req).build();
            String text = agent.text(req.getConversationId(), req.getMessage());
            return text;
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    @Override
    public Response<Image> image(ImageReq req) {
        try {
            ImageModel model = provider.image(req.getModelId());
            return model.generate(req.getPrompt().text());
        } catch (Exception e) {
            e.printStackTrace();
            throw new BusinessException("图片生成失败");
        }
    }
}