from datacloud.dev_utils.ydb.lib.core.ydb_lazy_driver import YdbLazyDriver
from datacloud.dev_utils.ydb.lib.core.ydb_lazy_session_pool import YdbLazySessionPool


__all__ = [
    'YdbManager',
    'YdbConnectionParams'
]


class YdbConnectionParams:
    __slots__ = ('endpoint', 'database', 'auth_token',
                 'tvm_client_id', 'tvm_secret', 'tvm_ydb_id')

    def __init__(self, endpoint, database, auth_token=None,
                 tvm_client_id=0, tvm_secret=None, tvm_ydb_id=0):
        assert auth_token or (tvm_client_id and tvm_secret and tvm_ydb_id), \
            'Please provide auth_token OR all tvm credentials'
        self.endpoint = endpoint
        self.database = database
        self.auth_token = auth_token
        self.tvm_client_id = int(tvm_client_id) if tvm_client_id else None
        self.tvm_secret = tvm_secret
        self.tvm_ydb_id = int(tvm_ydb_id) if tvm_ydb_id else None

        self.use_tvm = False
        if self.tvm_client_id and self.tvm_secret and self.tvm_ydb_id:
            self.use_tvm = True


class YdbManager(object):
    __slots__ = ('_driver', '_session_pool')

    def __init__(self, connection_params):
        self._driver = YdbLazyDriver(connection_params)
        self._session_pool = YdbLazySessionPool(self._driver)

    @property
    def pool(self):
        return self._session_pool

    @property
    def driver(self):
        return self._driver
