from typing import Optional

from pymilvus import MilvusClient as PyMilvusClient

from app.conf.app_config import MilvusConfig, app_config


class MilvusClientWrapper:
    """Milvus 客户端单例包装器"""
    _instance = None
    _initialized = False

    def __new__(cls, milvus_config: MilvusConfig = None):
        if cls._instance is None:
            cls._instance = super(MilvusClientWrapper, cls).__new__(cls)
        return cls._instance

    def __init__(self, milvus_config: MilvusConfig = None):
        if self._initialized:
            return

        if milvus_config is None:
            milvus_config = app_config.milvus

        self.milvus_config = milvus_config
        self.client: Optional[PyMilvusClient] = None
        self._initialized = True

    def init(self):
        if self.client is not None:
            return

        uri = f"http://{self.milvus_config.host}:{self.milvus_config.port}"
        self.client = PyMilvusClient(
            uri=uri,
            user=self.milvus_config.user,
            password=self.milvus_config.password
        )

    def close(self):
        if self.client is not None:
            self.client.close()
            self.client = None

    def has_collection(self, collection_name: str) -> bool:
        self._ensure_connection()
        return self.client.has_collection(collection_name)

    def drop_collection(self, collection_name: str):
        self._ensure_connection()
        if self.client.has_collection(collection_name):
            self.client.drop_collection(collection_name)

    def _ensure_connection(self):
        if self.client is None:
            self.init()


milvus_client = MilvusClientWrapper()