from sqlalchemy.orm import Session

from app.config import settings

from common_logging import get_logger

logger = get_logger(__name__)




def get_vector_store(
    db: Session, knowledge_base_id: int | None = None, tenant_id: int | None = None
):
    if settings.VECTOR_STORE_TYPE == "milvus":
        if settings.USE_MILVUS_LITE:
            from app.services.storage.milvus_basic import get_vector_store as get_milvus_basic_store

            logger.info(f"Using Milvus Lite vector store (knowledge base ID: {knowledge_base_id})")
            return get_milvus_basic_store(
                db=db,
                knowledge_base_id=knowledge_base_id,
                host=settings.MILVUS_HOST,
                port=settings.MILVUS_PORT,
                use_lite=True,
                tenant_id=tenant_id,
            )
        else:
            from app.services.storage.milvus import get_vector_store as get_milvus_store

            logger.info(
                f"Using Milvus vector store (knowledge base ID: {knowledge_base_id}, tenant ID: {tenant_id})"
            )
            return get_milvus_store(
                db=db,
                knowledge_base_id=knowledge_base_id,
                host=settings.MILVUS_HOST,
                port=settings.MILVUS_PORT,
                use_partition=True,
                tenant_id=tenant_id,
            )
    elif settings.VECTOR_STORE_TYPE == "pgvector":
        from app.services.storage.pgvector import get_vector_store as get_pgvector_store


        logger.info(f"Using pgvector vector store (knowledge base ID: {knowledge_base_id})")
        return get_pgvector_store(db, knowledge_base_id)
    else:
        raise ValueError(f"Unsupported vector store type: {settings.VECTOR_STORE_TYPE}")