from datacloud.dev_utils.ydb.lib.core.ydb_table import YdbTable
import ydb


class CryptaTable(YdbTable):
    def __init__(self, ydb_manager, database, table):
        super(CryptaTable, self).__init__(ydb_manager, database, table)

    def create(self):
        with self._init_session() as session:
            session.create_table(
                self.full_table_path,
                ydb.TableDescription()
                .with_column(ydb.Column('hashed_id', ydb.OptionalType(ydb.DataType.Uint64)))
                .with_column(ydb.Column('hashed_cid', ydb.OptionalType(ydb.DataType.Uint64)))
                .with_primary_key('hashed_id')
            )

    def get_one(self, record):
        query_params = {
            '$hashed_id': record.hashed_id
        }
        return self._get_one(query_params)

    def get(self, record):
        query_params = {
            '$hashed_cid': record.hashed_id
        }
        for rec in self._get(query_params):
            yield rec

    def get_multiple(self, id_values_list):
        # TODO: Raw request, refactor later
        with self._init_session() as session:
            query = self._select_multiple_request.format(
                database=self.database,
                table=self.table,
                items='(' + ', '.join(map(str, id_values_list)) + ')'
            )
            result_sets = session.transaction(ydb.SerializableReadWrite()).execute(
                query,
                commit_tx=True
            )
            for row in result_sets[0].rows:
                yield self.Record(**row)

    class Record(object):
        __slots__ = ('hashed_id', 'hashed_cid')

        def __init__(self, hashed_id, hashed_cid=None):
            self.hashed_id = hashed_id
            self.hashed_cid = hashed_cid

        def __str__(self):
            return '({hashed_id}, {hashed_cid})'.format(hashed_id=self.hashed_id, hashed_cid=self.hashed_cid)

        def __repr__(self):
            return self.__str__()

        def __eq__(self, other):
            return self.hashed_id == other.hashed_id and self.hashed_cid == other.hashed_cid

    _insert_request = """
        PRAGMA TablePathPrefix("{database}");

        DECLARE $records AS "List<Struct<
            hashed_id: Uint64,
            hashed_cid: Uint64>>";

        REPLACE INTO [{table}]
        SELECT
            hashed_id,
            hashed_cid
        FROM AS_TABLE($records);
    """

    _select_request = """
        PRAGMA TablePathPrefix("{database}");

        DECLARE $hashed_id AS Uint64;

        SELECT *
        FROM [{table}]
        WHERE hashed_id = $hashed_id;
    """

    _select_multiple_request = """
        PRAGMA TablePathPrefix("{database}");

        SELECT *
        FROM [{table}]
        WHERE hashed_id IN {items};
    """
