package tech.easyflow.ai.config;

import com.easyagents.engine.es.ElasticSearcher;
import com.easyagents.search.engine.service.DocumentSearcher;
import com.easyagents.store.milvus.MilvusVectorStore;
import org.springframework.boot.actuate.health.Health;
import org.springframework.boot.actuate.health.HealthIndicator;
import org.springframework.stereotype.Component;
import tech.easyflow.ai.rag.KeywordEngineType;
import tech.easyflow.ai.support.DocumentStoreLifecycleSupport;
import tech.easyflow.common.util.SpringContextUtil;
import tech.easyflow.common.util.StringUtil;

import java.io.File;

/**
 * RAG 依赖中间件健康检查。
 */
public class RagHealthIndicator {

    /**
     * Milvus 健康检查。
     */
    @Component("ragMilvusHealthIndicator")
    public static class RagMilvusHealthIndicator extends CachedHealthIndicatorSupport implements HealthIndicator {

        private final AiMilvusConfig aiMilvusConfig;

        /**
         * 创建 Milvus 健康检查器。
         *
         * @param aiMilvusConfig Milvus 配置
         * @param healthProperties RAG 健康检查配置
         */
        public RagMilvusHealthIndicator(AiMilvusConfig aiMilvusConfig, RagHealthProperties healthProperties) {
            super(healthProperties);
            this.aiMilvusConfig = aiMilvusConfig;
        }

        /**
         * 检查 Milvus 是否可连接。
         *
         * @return 健康状态
         */
        @Override
        public Health health() {
            return cachedHealth();
        }

        @Override
        protected Health doHealthCheck() {
            MilvusVectorStore vectorStore = null;
            try {
                vectorStore = new MilvusVectorStore(
                    aiMilvusConfig.copyForCollection("__rag_health_probe__")
                );
                if (vectorStore.checkAvailable()) {
                    return Health.up().withDetail("uri", aiMilvusConfig.getUri()).build();
                }
                return Health.down().withDetail("uri", aiMilvusConfig.getUri()).build();
            } catch (Exception e) {
                return Health.down(e).withDetail("uri", aiMilvusConfig.getUri()).build();
            } finally {
                DocumentStoreLifecycleSupport.closeQuietly(vectorStore);
            }
        }
    }

    /**
     * 关键词检索健康检查。
     */
    @Component("ragKeywordSearchHealthIndicator")
    public static class RagKeywordSearchHealthIndicator extends CachedHealthIndicatorSupport implements HealthIndicator {

        private final SearcherFactory searcherFactory;
        private final AiLuceneConfig aiLuceneConfig;

        /**
         * 创建关键词检索健康检查器。
         *
         * @param searcherFactory 检索器工厂
         * @param aiLuceneConfig Lucene 配置
         * @param healthProperties RAG 健康检查配置
         */
        public RagKeywordSearchHealthIndicator(SearcherFactory searcherFactory,
                                               AiLuceneConfig aiLuceneConfig,
                                               RagHealthProperties healthProperties) {
            super(healthProperties);
            this.searcherFactory = searcherFactory;
            this.aiLuceneConfig = aiLuceneConfig;
        }

        /**
         * 检查当前关键词检索引擎是否可用。
         *
         * @return 健康状态
         */
        @Override
        public Health health() {
            return cachedHealth();
        }

        @Override
        protected Health doHealthCheck() {
            KeywordEngineType engineType = KeywordEngineType.from(SpringContextUtil.getProperty("rag.engine", "ES"));
            if (engineType == KeywordEngineType.LUCENE) {
                return checkLuceneDirectory(engineType);
            }
            DocumentSearcher searcher = searcherFactory.getSearcher();
            if (searcher instanceof ElasticSearcher elasticSearcher && elasticSearcher.checkAvailable()) {
                return Health.up().withDetail("engine", engineType.name()).build();
            }
            return Health.down().withDetail("engine", engineType.name()).build();
        }

        private Health checkLuceneDirectory(KeywordEngineType engineType) {
            String indexDirPath = aiLuceneConfig.getIndexDirPath();
            if (StringUtil.noText(indexDirPath)) {
                return Health.down()
                    .withDetail("engine", engineType.name())
                    .withDetail("reason", "Lucene 索引目录未配置")
                    .build();
            }
            File indexDir = new File(indexDirPath);
            if (indexDir.exists() && indexDir.isDirectory() && indexDir.canRead() && indexDir.canWrite()) {
                return Health.up()
                    .withDetail("engine", engineType.name())
                    .withDetail("indexDir", indexDirPath)
                    .build();
            }
            return Health.down()
                .withDetail("engine", engineType.name())
                .withDetail("indexDir", indexDirPath)
                .withDetail("reason", "Lucene 索引目录不可读写")
                .build();
        }
    }
}
