package com.ruoyi.common.utils.milvus;

import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.service.collection.request.LoadCollectionReq;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.util.*;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

@Component
public class MilvusConnectionPool {

    @Value("${cmc.milvusService.url}")
    private String milvusServiceUrl;

    private static final int POOL_SIZE = 5;
    private final List<MilvusClientV2> clientPool = new ArrayList<>();
    private final BlockingQueue<MilvusClientV2> availableClients = new LinkedBlockingQueue<>();
    private final Set<String> loadedCollections = ConcurrentHashMap.newKeySet();
    private final ReadWriteLock collectionLock = new ReentrantReadWriteLock();
    private volatile boolean initialized = false;

    @PostConstruct
    public void init() {
        if (milvusServiceUrl == null || milvusServiceUrl.isEmpty()) {
            throw new IllegalArgumentException("milvusServiceUrl 配置不能为空");
        }

        for (int i = 0; i < POOL_SIZE; i++) {
            // MilvusClientV2 client = new MilvusClientV2(
            //         ConnectConfig.builder()
            //                 .uri(milvusServiceUrl)
            //                 .build());
            // clientPool.add(client);
            // availableClients.offer(client);
        }
        initialized = true;
        System.out.println("Milvus连接池初始化完成，连接数: " + POOL_SIZE);
    }

    @PreDestroy
    public void destroy() {
        for (MilvusClientV2 client : clientPool) {
            try {
                if (client != null) {
                    client.close();
                }
            } catch (Exception e) {
                System.err.println("关闭Milvus连接时出错: " + e.getMessage());
            }
        }
        clientPool.clear();
        availableClients.clear();
        System.out.println("Milvus连接池已关闭");
    }

    public MilvusClientV2 getClient() throws InterruptedException {
        MilvusClientV2 client = availableClients.poll(5, TimeUnit.SECONDS);
        if (client == null) {
            throw new RuntimeException("获取Milvus连接超时");
        }
        return client;
    }

    public void returnClient(MilvusClientV2 client) {
        if (client != null) {
            availableClients.offer(client);
        }
    }

    public void ensureCollectionLoaded(MilvusClientV2 client, String collectionName) {
        collectionLock.readLock().lock();
        try {
            if (loadedCollections.contains(collectionName)) {
                return;
            }
        } finally {
            collectionLock.readLock().unlock();
        }

        collectionLock.writeLock().lock();
        try {
            if (!loadedCollections.contains(collectionName)) {
                LoadCollectionReq loadCollectionReq = LoadCollectionReq.builder()
                        .collectionName(collectionName)
                        .build();
                client.loadCollection(loadCollectionReq);
                loadedCollections.add(collectionName);
                System.out.println("集合已加载: " + collectionName);
            }
        } finally {
            collectionLock.writeLock().unlock();
        }
    }

    public boolean isInitialized() {
        return initialized;
    }

    public Set<String> getLoadedCollections() {
        return new HashSet<>(loadedCollections);
    }
}