package tech.easyflow.ai.service.impl;

import com.easyagents.rag.retrieval.RetrievalMode;
import tech.easyflow.ai.dto.BotKnowledgeBindingRequest;
import tech.easyflow.ai.entity.BotDocumentCollection;
import tech.easyflow.ai.mapper.BotDocumentCollectionMapper;
import tech.easyflow.ai.service.BotDocumentCollectionService;
import com.mybatisflex.spring.service.impl.ServiceImpl;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import tech.easyflow.common.cache.RedisLockExecutor;

import javax.annotation.Resource;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.time.Duration;
import com.mybatisflex.core.query.QueryWrapper;

/**
 *  服务层实现。
 *
 * @author michael
 * @since 2024-08-28
 */
@Service
public class BotDocumentCollectionServiceImpl extends ServiceImpl<BotDocumentCollectionMapper, BotDocumentCollection> implements BotDocumentCollectionService {

    private static final String BOT_BINDING_LOCK_KEY_PREFIX = "easyflow:lock:bot:binding:";
    private static final Duration LOCK_WAIT_TIMEOUT = Duration.ofSeconds(2);
    private static final Duration LOCK_LEASE_TIMEOUT = Duration.ofSeconds(10);

    @Resource
    private RedisLockExecutor redisLockExecutor;

    @Override
    public List<BotDocumentCollection> listByBotId(BigInteger botId) {

        QueryWrapper queryWrapper = QueryWrapper.create();
        queryWrapper.eq(BotDocumentCollection::getBotId,botId);

        return list(queryWrapper);
    }

    @Override
    @Transactional
    public void saveBotAndKnowledge(BigInteger botId, List<BotKnowledgeBindingRequest> knowledgeBindings) {
        redisLockExecutor.executeWithLock(BOT_BINDING_LOCK_KEY_PREFIX + botId, LOCK_WAIT_TIMEOUT, LOCK_LEASE_TIMEOUT, () -> {
            this.remove(QueryWrapper.create().eq(BotDocumentCollection::getBotId, botId));
            Map<BigInteger, RetrievalMode> bindingMap = new LinkedHashMap<>();
            if (knowledgeBindings != null) {
                for (BotKnowledgeBindingRequest binding : knowledgeBindings) {
                    if (binding == null || binding.getKnowledgeId() == null) {
                        continue;
                    }
                    bindingMap.put(binding.getKnowledgeId(), binding.resolveRetrievalMode());
                }
            }
            if (bindingMap.isEmpty()) {
                return;
            }

            List<BotDocumentCollection> list = new ArrayList<>(bindingMap.size());
            for (Map.Entry<BigInteger, RetrievalMode> entry : bindingMap.entrySet()) {
                BotDocumentCollection botDocumentCollection = new BotDocumentCollection();
                botDocumentCollection.setBotId(botId);
                botDocumentCollection.setDocumentCollectionId(entry.getKey());
                Map<String, Object> options = new HashMap<>();
                options.put(BotDocumentCollection.OPTION_KEY_RETRIEVAL_MODE, entry.getValue().name());
                botDocumentCollection.setOptions(options);
                list.add(botDocumentCollection);
            }
            this.saveBatch(list);
        });
    }
}
