#用户请求限制规则配置

// 敏感词过滤规则
@Slf4j
@Component("sensitive")
@RequiredArgsConstructor
public class SensitiveWordsRule extends NodeComponent {

    private final AiChatRecordMapper aiChatRecordMapper;
    private final AiKnowledgeProperties aiKnowledgeProperties;

    @Override
    public void process() {
        ChatMessageContext messageContext = this.getContextBean(ChatMessageContext.class);
        AiDatasetEntity aiDataset = messageContext.getAiDataset();

        // 判断是否开启了敏感词过滤
        if (YesNoEnum.NO.getCode().equals(aiDataset.getSensitiveFlag())) {
            return;
        }

        // 敏感词过滤
        SensitiveWordBs sensitiveWordBs = SpringUtil.getBean(SensitiveWordBs.class);
        if (sensitiveWordBs.contains(messageContext.getChatMessageDTO().getContent())) {
            Long messageKey = messageContext.getChatMessageDTO().getMessageKey();
            AiChatRecordEntity aiChatRecordEntity = new AiChatRecordEntity();
            messageContext.setErrorMessage(aiDataset.getSensitiveMsg());
            setIsEnd(true);  // 设置错误消息并终止流程
        }
    }
}

// 用户维度风控规则
@Slf4j
@Component("user")
@RequiredArgsConstructor
public class UserRiskControlRule extends NodeComponent {

    private final AiKnowledgeProperties aiKnowledgeProperties;
    private final AiBillService aiBillService;

    @Override
    public void process() {
        ChatMessageContext messageContext = this.getContextBean(ChatMessageContext.class);
        String risked = riskControl(messageContext.getUsername());
        if (StrUtil.isNotBlank(risked)) {
            messageContext.setErrorMessage(risked);
            setIsEnd(true);
        }
    }

    private String riskControl(String username) {
        // 统计用户从昨天到现在的调用次数
        long count = aiBillService.count(Wrappers.<AiBillEntity>lambdaQuery()
            .eq(StrUtil.isNotBlank(username), AiBillEntity::getCreateBy, username)
            .ge(AiBillEntity::getCreateTime, DateUtil.yesterday()));

        log.info("username:{}, count:{}", username, count);
        if (count > riskControl.getTimes()) {
            return riskControl.getErrorMsg();
        }
        return null;
    }
}

// 白名单判断规则
@Slf4j
@Component("isNoLimit")
@RequiredArgsConstructor
public class IsNoLimitRule extends NodeBooleanComponent {

    private final AiKnowledgeProperties aiKnowledgeProperties;

    @Override
    public boolean processBoolean() throws Exception {
        // 如果白名单为空，所有用户都需要风控
        if (CollUtil.isEmpty(aiKnowledgeProperties.getRiskControl().getNoLimitUsernames())) {
            return false;
        }

        String username = SecurityUtils.getUser().getUsername();
        return aiKnowledgeProperties.getRiskControl().getNoLimitUsernames().contains(username);
    }
}

// 聊天类型选择规则
@Slf4j
@Component("selectChatType")
public class ChatTypeSelectRule extends NodeSwitchComponent {

    @Override
    public String processSwitch() throws Exception {
        ChatMessageContext messageContext = this.getContextBean(ChatMessageContext.class);
        ChatMessageDTO lastUserMessage = messageContext.getChatMessageDTO();

        // 根据数据集ID选择对应的聊天类型
        ChatTypeEnums chatTypeEnums = ChatTypeEnums.fromCode(lastUserMessage.getDatasetId());
        return Objects.requireNonNullElse(chatTypeEnums, ChatTypeEnums.VECTOR_CHAT).getType();
    }
}

// 自定义内容质量风控规则
@Slf4j
@Component("contentQuality")
@RequiredArgsConstructor
public class ContentQualityRule extends NodeComponent {

    private final YourCustomService customService;

    @Override
    public void process() {
        ChatMessageContext messageContext = this.getContextBean(ChatMessageContext.class);
        String content = messageContext.getChatMessageDTO().getContent();

        boolean isLowQuality = checkContentQuality(content);

        if (isLowQuality) {
            log.warn("内容质量不达标，用户：{}", messageContext.getUsername());
            messageContext.setErrorMessage("输入内容质量不符合要求，请重新输入");
            setIsEnd(true);
        }
    }

    private boolean checkContentQuality(String content) {
        // 自定义质量检查逻辑
        if (content.length() < 5) {
            return true;
        }
        return customService.scoreContent(content) < 0.6;
    }
}

// 自定义聊天规则实现
@Slf4j
@Component("customChat")
@RequiredArgsConstructor
public class CustomChatRule implements ChatRule {

    private final ModelProvider modelProvider;

    @Override
    public Flux<AiMessageResultDTO> process(ChatMessageDTO chatMessageDTO) {
        AiStreamAssistantService streamAssistantService = modelProvider
            .getAiStreamAssistant(chatMessageDTO.getModelName())
            .getValue();

        return streamAssistantService
            .chat(chatMessageDTO.getConversationId(), chatMessageDTO.getContent())
            .map(AiMessageResultDTO::new);
    }
}

public enum ChatTypeEnums {
    VECTOR_CHAT("vectorChat", "向量聊天"),
    SIMPLE_CHAT("simpleChat", "简单聊天"),
    CUSTOM_CHAT("customChat", "自定义聊天"),  // 新增
    // ... 其他类型
}