package com.xly.milvus.config;

import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.service.collection.request.HasCollectionReq;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.ApplicationArguments;
import org.springframework.boot.ApplicationRunner;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

import java.util.List;

@Slf4j
@Component
@Order(1)
public class MilvusStartupValidator implements ApplicationRunner {

    @Autowired
    private MilvusClientV2 milvusClient;

    @Value("${milvus.validation.collections:}")
    private List<String> validationCollections;

    @Value("${milvus.validation.timeout:5000}")
    private long validationTimeout;

    @Override
    public void run(ApplicationArguments args) throws Exception {
        log.info("开始Milvus启动验证...");

        try {
            // 1. 验证客户端就绪状态
            boolean isReady = milvusClient.clientIsReady();
            log.info("客户端就绪状态: {}", isReady);

            // 2. 获取服务器版本
            String serverVersion = milvusClient.getServerVersion();
            log.info("Milvus服务器版本: {}", serverVersion);

            // 3. 列出所有数据库
            var databases = milvusClient.listDatabases();
            log.info("可用数据库: {}", databases.getDatabaseNames());

            // 4. 检查当前使用的数据库
            String currentDb = milvusClient.currentUsedDatabase();
            log.info("当前数据库: {}", currentDb);

            // 5. 健康检查
            var healthCheck = milvusClient.checkHealth();
            log.info("健康状态: {}", healthCheck.getIsHealthy());
            if (healthCheck.getQuotaStates() != null && !healthCheck.getQuotaStates().isEmpty()) {
                log.info("配额状态: {}", healthCheck.getQuotaStates());
            }

            // 6. 验证指定的集合（可选）
            if (validationCollections != null && !validationCollections.isEmpty()) {
                for (String collectionName : validationCollections) {
                    boolean exists = milvusClient.hasCollection(
                            HasCollectionReq.builder()
                                    .collectionName(collectionName)
                                    .build()
                    );
                    log.info("集合 '{}' 存在: {}", collectionName, exists);
                }
            }

            log.info("Milvus启动验证完成");

        } catch (Exception e) {
            log.error("Milvus启动验证失败: {}", e.getMessage(), e);
            // 可以根据配置决定是否抛出异常
            // throw new RuntimeException("Milvus验证失败", e);
        }
    }
}