package com.xly.milvus.config;

import io.milvus.v2.client.MilvusClientV2;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.actuate.health.Health;
import org.springframework.boot.actuate.health.HealthIndicator;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.Map;

@Component
public class MilvusHealthIndicator implements HealthIndicator {

    @Autowired(required = false)
    private MilvusClientV2 milvusClient;

    @Override
    public Health health() {
        if (milvusClient == null) {
            return Health.down()
                    .withDetail("error", "Milvus客户端未初始化")
                    .build();
        }

        try {
            // 检查客户端是否就绪（源码中存在clientIsReady()）
            boolean isReady = milvusClient.clientIsReady();

            if (!isReady) {
                return Health.down()
                        .withDetail("error", "客户端未就绪")
                        .build();
            }

            long startTime = System.currentTimeMillis();

            // 执行健康检查 - 使用源码中存在的方法
            String serverVersion = milvusClient.getServerVersion();
            var databases = milvusClient.listDatabases();
            var healthCheck = milvusClient.checkHealth();

            long responseTime = System.currentTimeMillis() - startTime;

            Map<String, Object> details = new HashMap<>();
            details.put("serverVersion", serverVersion);
            details.put("database", milvusClient.currentUsedDatabase());
            details.put("databases", databases.getDatabaseNames());
            details.put("healthStatus", healthCheck.getIsHealthy());
            details.put("responseTime", responseTime + "ms");
            details.put("clientReady", isReady);

            if (healthCheck.getQuotaStates() != null) {
                details.put("quotaStates", healthCheck.getQuotaStates());
            }

            return Health.up()
                    .withDetails(details)
                    .build();

        } catch (Exception e) {
            return Health.down()
                    .withDetail("error", e.getMessage())
                    .withDetail("clientReady", milvusClient.clientIsReady())
                    .build();
        }
    }
}