import ydb
from kikimr.public.sdk.python import tvm
from concurrent.futures import TimeoutError
from threading import Lock


__all__ = [
    'YdbLazyDriver'
]


class YdbLazyDriver(object):
    def __init__(self, connection_params, timeout=5):
        self._connection_params = connection_params
        self._driver = None
        self._driver_lock = Lock()
        self._timeout = timeout

    def _get_driver(self):
        if self._connection_params.use_tvm:
            print('[AUTH]: Use TVM')
            credentials = tvm.TvmCredentialsProvider(
                self_client_id=int(self._connection_params.tvm_client_id),
                self_secret=self._connection_params.tvm_secret,
                destination_alias='ydb',
                dsts={'ydb': int(self._connection_params.tvm_ydb_id)})
        else:
            print('[AUTH]: Use Token')
            credentials = ydb.AuthTokenCredentials(
                self._connection_params.auth_token)
        driver = ydb.Driver(
            ydb.ConnectionParams(
                self._connection_params.endpoint,
                database=self._connection_params.database,
                credentials=credentials
                # auth_token=self._connection_params.auth_token
            )
        )
        try:
            driver.wait(timeout=self._timeout)
        except TimeoutError:
            raise RuntimeError('YdbLazyDriver - Connection Error.')
        return driver

    @property
    def driver(self):
        if self._driver is None:
            print('Create lazy driver')
            self._driver = self._get_driver()
        return self._driver
