import yt.wrapper as yt_wrapper
import ydb
from datacloud.dev_utils.logging.logger import get_basic_logger
from datacloud.dev_utils.ydb.lib.core import utils as ydb_utils

logger = get_basic_logger(__name__)


class Deployer:
    def __init__(self, ydb_manager, yt_client=None):
        self._manager = ydb_manager
        self._yt_client = yt_wrapper  # TODO: Change to get_yt_client, was used before Arcadia

    def create_ydb_table(self, db_info, harvesting):
        logger.info('[DEPLOY] Create table %s', harvesting)
        table_folder = db_info.get_ydb_table_folder(harvesting)
        ydb_utils.create_folder(
            self._manager.driver.driver, db_info.ydb_database, table_folder)
        self.get_ydb_table(db_info, harvesting).create()

    def ydb_records(self, db_info, harvesting):
        table_path = db_info.get_ydb_table_path(harvesting)
        with self._manager.pool.acquire() as session:
            query = self._select_all_request.format(
                database=db_info.ydb_database, table=table_path)
            prepared_query = session.prepare(query)
            response = session.transaction(ydb.SerializableReadWrite()).execute(
                prepared_query, {}, commit_tx=True)
            for row in response[0].rows:
                yield row

    def get_ydb_table(self, db_info, harvesting):
        table = db_info.get_ydb_table_path(harvesting)
        return harvesting.ydb_cls(self._manager, db_info.ydb_database, table)

    def ydb_upload(self, db_info, harvesting, records):
        logger.info('[DEPLOY] Fill data into %s', harvesting)
        records = list(map(lambda it: harvesting.ydb_cls.Record(**it), records))
        self.get_ydb_table(db_info, harvesting).insert(records)

    def yt_records(self, db_info, harvesting):
        assert db_info.is_yt_available, 'Pleas enable YT ops in DBInfo'
        table_path = self._get_yt_table_path(harvesting)
        return self._yt_client.read_table(table_path)

    def yt_to_ydb(self, db_info, harvesting):
        self._check_yt_presence()
        raise 'Yt to Ydb is not implemented'

    def ydb_to_yt(self, db_info, harvesting):
        self._check_yt_presence()
        raise 'Ydb to Yt is not implemented'

    def _infer_yt_schema():
        raise 'Ydb '

    def create_yt_table(self, db_info, harvesting):
        self._check_yt_presence()
        raise 'Create YT table is not implemented'

    def deploy(self, recipe):
        for harvesting in recipe.harvesting_group:
            self.create_ydb_table(recipe.db_info, harvesting)

    def fill(self, recipe, data_source=None):
        data_source = data_source or recipe.data_source
        assert data_source is not None, 'Please, provide data source'
        for harvesting in recipe.harvesting_group:
            records = data_source.get_data(harvesting)
            if records:
                self.ydb_upload(recipe.db_info, harvesting, records)
            else:
                logger.warn('No data for %s', harvesting)

    def dump(self, recipe, data_source):
        for harvesting in recipe.harvesting_group:
            records = list(self.ydb_records(recipe.db_info, harvesting))
            data_source.set_data(harvesting, records)
        return data_source

    _select_all_request = """
        PRAGMA TablePathPrefix("{database}");
        SELECT * FROM [{table}];
    """
