
package tech.easyflow.ai.service.impl;

import cn.hutool.core.util.StrUtil;
import com.alicp.jetcache.Cache;
import com.easyagents.core.document.Document;
import com.easyagents.core.model.embedding.EmbeddingModel;
import com.easyagents.core.model.rerank.RerankModel;
import com.easyagents.core.store.VectorData;
import com.mybatisflex.core.query.QueryWrapper;
import com.mybatisflex.core.util.StringUtil;
import com.mybatisflex.spring.service.impl.ServiceImpl;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.CollectionUtils;
import tech.easyflow.ai.entity.Model;
import tech.easyflow.ai.entity.ModelProvider;
import tech.easyflow.ai.mapper.ModelMapper;
import tech.easyflow.ai.service.ModelProviderService;
import tech.easyflow.ai.service.ModelService;
import tech.easyflow.ai.service.verification.ChatModelConnectivityVerifier;
import tech.easyflow.common.tree.Tree;
import tech.easyflow.common.util.SqlOperatorsUtil;
import tech.easyflow.common.util.SqlUtil;
import tech.easyflow.common.web.exceptions.BusinessException;

import javax.annotation.Resource;
import java.math.BigInteger;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 服务层实现。
 *
 * @author michael
 * @since 2024-08-23
 */
@Service
public class ModelServiceImpl extends ServiceImpl<ModelMapper, Model> implements ModelService {

    @Autowired
    ModelMapper modelMapper;

    @Autowired
    ModelProviderService modelProviderService;

    @Resource
    private Cache<String, Object> cache;

    /** 与智能体运行时同链路的 Chat Model 连通性验证器。 */
    @Autowired(required = false)
    private ChatModelConnectivityVerifier chatModelConnectivityVerifier;


    @Override
    public boolean addAiLlm(Model entity) {
        int insert = modelMapper.insert(entity);
        if (insert <= 0) {
            return false;
        }
        return true;
    }

    private static final Logger log = LoggerFactory.getLogger(ModelServiceImpl.class);

    @Override
    public Map<String, Object> verifyModelConfig(Model model) {
        String modelType = model.getModelType();
        Map<String, Object> resMap = new HashMap<>();
        // 走聊天验证逻辑
        if (Model.MODEL_TYPES[0].equals(modelType)) {
            return verifyChatLlm(model);
        }
        // 走向量化验证逻辑
        if (Model.MODEL_TYPES[1].equals(modelType)) {
            int dimension = verifyEmbedLlm(model);
            resMap.put("dimension", dimension);
            return resMap;
        }
        // 走重排验证逻辑
        if (Model.MODEL_TYPES[2].equals(modelType)) {
            verifyRerankLlm(model);
            return null;

        }

        // 以上不满足，视为验证失败
        throw new BusinessException("校验失败！");

    }

    @Override
    public Map<String, Map<String, List<Model>>> getList(Model entity) {
        Map<String, Map<String, List<Model>>> result = new HashMap<>();

        QueryWrapper queryWrapper = new QueryWrapper()
                .eq(Model::getProviderId, entity.getProviderId());
        List<Model> totalList = modelMapper.selectListWithRelationsByQuery(queryWrapper);
        for (String modelType : Model.MODEL_TYPES) {
            Map<String, List<Model>> groupMap = groupLlmByGroupName(totalList, modelType);
            if (!CollectionUtils.isEmpty(groupMap)) {
                result.put(modelType, groupMap);
            }
        }

        return result;
    }

    private Map<String, List<Model>> groupLlmByGroupName(List<Model> totalList, String targetModelType) {
        if (CollectionUtils.isEmpty(totalList)) {
            return Collections.emptyMap();
        }

        return totalList.stream()
                .filter(aiLlm -> targetModelType.equals(aiLlm.getModelType())
                        && aiLlm.getGroupName() != null)
                .collect(Collectors.groupingBy(Model::getGroupName));
    }


    private void verifyRerankLlm(Model model) {
        RerankModel rerankModel = model.toRerankModel();
        List<Document> documents = new ArrayList<>();
        documents.add(Document.of("Paris is the capital of France."));
        documents.add(Document.of("London is the capital of England."));
        documents.add(Document.of("Tokyo is the capital of Japan."));
        documents.add(Document.of("Beijing is the capital of China."));
        documents.add(Document.of("Washington, D.C. is the capital of the United States."));
        documents.add(Document.of("Moscow is the capital of Russia."));
        try {
            List<Document> rerank = rerankModel.rerank("What is the capital of France?", documents);
            if (rerank == null || rerank.isEmpty()) {
                throw new BusinessException("校验未通过，请前往后端日志查看详情！");
            }
        } catch (Exception e) {
            log.error("校验失败：{}", e.getMessage());
            throw new BusinessException(e.getMessage());
        }
    }

    private int verifyEmbedLlm(Model model) {
        try {
            EmbeddingModel embeddingModel = model.toEmbeddingModel();
            VectorData vectorData = embeddingModel.embed("这是一条校验模型配置的文本");
            if (vectorData.getVector() == null) {
                throw new BusinessException("校验未通过，请前往后端日志查看详情！");
            }
            log.info("取到向量数据，校验结果通过");
            return vectorData.getVector().length;
        } catch (Exception e) {
            log.error("模型配置校验失败:{}", e.getMessage());
            throw new BusinessException(e.getMessage());
        }
    }

    /**
     * 使用智能体同链路验证 Chat Model 或 VLM。
     *
     * @param model 已补齐供应商默认配置的模型
     * @return 结构化双阶段验证结果
     */
    private Map<String, Object> verifyChatLlm(Model model) {
        if (chatModelConnectivityVerifier == null) {
            throw new BusinessException("Agent 模型连通性验证组件未加载");
        }
        return chatModelConnectivityVerifier.verify(model).toMap();
    }

    @Override
    public void removeByEntity(Model entity) {
        QueryWrapper queryWrapper = QueryWrapper.create().eq(Model::getProviderId, entity.getProviderId()).eq(Model::getGroupName, entity.getGroupName());
        modelMapper.deleteByQuery(queryWrapper);
    }

    @Override
    public Model getModelInstance(BigInteger modelId) {
        if (modelId == null) {
            throw new BusinessException("模型ID不能为空");
        }
        Model model = modelMapper.selectOneWithRelationsById(modelId);
        return fillProviderDefaults(model);
    }

    @Override
    public Model getModelInstanceByInvokeCode(String invokeCode) {
        if (StrUtil.isBlank(invokeCode)) {
            throw new BusinessException("invokeCode不能为空");
        }
        QueryWrapper queryWrapper = QueryWrapper.create().eq(Model::getInvokeCode, invokeCode.trim());
        Model model = modelMapper.selectOneWithRelationsByQuery(queryWrapper);
        return fillProviderDefaults(model);
    }

    @Override
    public void validateForSaveOrUpdate(Model entity, boolean isSave) {
        if (entity == null) {
            throw new BusinessException("模型配置不能为空");
        }
        if (entity.getPublishEnabled() == null) {
            entity.setPublishEnabled(Boolean.FALSE);
        }

        String originalInvokeCode = StrUtil.trim(entity.getInvokeCode());
        String invokeCode = originalInvokeCode;
        boolean autoGeneratedInvokeCode = StrUtil.isBlank(invokeCode);
        if (autoGeneratedInvokeCode) {
            invokeCode = buildDefaultInvokeCode(entity.getModelName());
        }

        if (Boolean.TRUE.equals(entity.getPublishEnabled())) {
            if (StrUtil.isBlank(invokeCode)) {
                throw new BusinessException("开启 API 调用前，请先配置 invokeCode");
            }
            if (!Model.MODEL_TYPES[0].equals(entity.getModelType())) {
                throw new BusinessException("只有聊天模型支持开启 API 调用");
            }
        }

        if (StrUtil.isBlank(invokeCode)) {
            entity.setInvokeCode(null);
            return;
        }

        QueryWrapper queryWrapper = QueryWrapper.create().eq(Model::getInvokeCode, invokeCode);
        if (!isSave && entity.getId() != null) {
            queryWrapper.ne(Model::getId, entity.getId());
        }
        boolean duplicated = modelMapper.selectCountByQuery(queryWrapper) > 0;
        if (duplicated && autoGeneratedInvokeCode && !Boolean.TRUE.equals(entity.getPublishEnabled())) {
            entity.setInvokeCode(null);
            return;
        }
        if (duplicated) {
            throw new BusinessException("invokeCode 已存在，请更换后重试");
        }
        entity.setInvokeCode(invokeCode);
    }

    @Override
    public List<Model> listInvokeModels() {
        QueryWrapper queryWrapper = QueryWrapper.create().eq(Model::getModelType, Model.MODEL_TYPES[0]);
        return modelMapper.selectListWithRelationsByQuery(queryWrapper).stream()
                .map(this::fillProviderDefaults)
                .collect(Collectors.toList());
    }

    @Override
    public List<Model> listSelectableModels(Model entity, Boolean asTree, String sortKey, String sortType) {
        QueryWrapper queryWrapper = QueryWrapper.create(
                entity,
                entity == null ? com.mybatisflex.core.query.SqlOperators.empty() : SqlOperatorsUtil.build(entity.getClass())
        );
        queryWrapper.orderBy(buildOrderBy(sortKey, sortType));
        List<Model> list = Tree.tryToTree(modelMapper.selectListWithRelationsByQuery(queryWrapper), asTree);
        list.forEach(this::decorateModelTitle);
        return list;
    }

    @Override
    public Model updateInvokeConfig(BigInteger id, String invokeCode, Boolean publishEnabled) {
        Model existing = getModelInstance(id);
        if (existing == null) {
            throw new BusinessException("模型不存在");
        }
        if (!Model.MODEL_TYPES[0].equals(existing.getModelType())) {
            throw new BusinessException("只有聊天模型支持统一网关配置");
        }

        existing.setInvokeCode(invokeCode);
        existing.setPublishEnabled(Boolean.TRUE.equals(publishEnabled));
        validateForSaveOrUpdate(existing, false);
        modelMapper.update(existing);
        return getModelInstance(id);
    }

    @Override
    @Transactional(rollbackFor = Exception.class)
    public List<Model> batchUpdateInvokePublishStatus(List<BigInteger> ids, Boolean publishEnabled) {
        if (CollectionUtils.isEmpty(ids)) {
            throw new BusinessException("请选择要操作的模型");
        }

        List<BigInteger> uniqueIds = new ArrayList<>(new LinkedHashSet<>(ids));
        List<Model> updatedModels = new ArrayList<>(uniqueIds.size());
        for (BigInteger id : uniqueIds) {
            updatedModels.add(updateInvokeConfig(id, null, publishEnabled));
        }
        return updatedModels;
    }

    private void decorateModelTitle(Model model) {
        if (model == null) {
            return;
        }
        String providerName = Optional.ofNullable(model.getModelProvider())
                .map(ModelProvider::getProviderName)
                .orElse("-");
        model.setTitle(providerName + "/" + model.getTitle());
    }

    private String buildOrderBy(String sortKey, String sortType) {
        sortKey = StringUtil.camelToUnderline(sortKey);
        return SqlUtil.buildOrderBy(sortKey, sortType, "id desc");
    }

    private String buildDefaultInvokeCode(String modelName) {
        if (StrUtil.isBlank(modelName)) {
            return null;
        }
        String normalized = modelName.trim()
                .replaceAll("[^A-Za-z0-9._:-]+", "-")
                .replaceAll("-{2,}", "-")
                .replaceFirst("^[^A-Za-z0-9]+", "");
        if (StrUtil.isBlank(normalized)) {
            return null;
        }
        if (normalized.length() > 128) {
            normalized = normalized.substring(0, 128);
        }
        if (normalized.length() == 1) {
            normalized = normalized + "-model";
        }
        return normalized;
    }

    private Model fillProviderDefaults(Model model) {
        if (model == null) {
            return null;
        }
        ModelProvider modelProvider = model.getModelProvider();
        model.setModelProvider(modelProvider);
        if (StrUtil.isBlank(model.getApiKey())) {
            model.setApiKey(modelProvider.getApiKey());
        }
        if (StrUtil.isBlank(model.getEndpoint())) {
            model.setEndpoint(modelProvider.getEndpoint());
        }

        // 请求路径为空，从modelProvider中获取
        if (StrUtil.isBlank(model.getRequestPath())) {
            // 模型类型为chatModel
            if (model.getModelType().equals(Model.MODEL_TYPES[0])) {
                model.setRequestPath(modelProvider.getChatPath());
                // 模型类型为embeddingModel
            } else if (model.getModelType().equals(Model.MODEL_TYPES[1])) {
                model.setRequestPath(modelProvider.getEmbedPath());
                // 模型类型为rerankModel
            } else if (model.getModelType().equals(Model.MODEL_TYPES[2])) {
                model.setRequestPath(modelProvider.getRerankPath());
            }
        }

        return model;
    }
}
