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


class PartnerTokensTable(YdbTable):
    def __init__(self, ydb_manager, database, table):
        super(PartnerTokensTable, 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('token', ydb.OptionalType(ydb.DataType.String)))
                .with_column(ydb.Column('partner_id', ydb.OptionalType(ydb.DataType.String)))
                .with_primary_key('token')
            )

    # TODO: Change order in Recrodd to token, partner.
    # TODO: Use record instead token
    def get_one(self, token):
        query_params = {
            '$token': token
        }
        return self._get_one(query_params)

    # def get(self, token):
    #     query_params = {
    #         '$token': token
    #     }
    #     for rec in self._get(query_params):
    #         yield rec

    class Record(YdbTable.Record):
        __slots__ = ('partner_id', 'token')

        def __init__(self, partner_id, token):
            self.token = token
            self.partner_id = partner_id

        def __str__(self):
            return '({}: {})'.format(self.partner_id, self.token)

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

        def __eq__(self, other):
            return self.token == other.token and self.partner_id == other.partner_id

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

        DECLARE $records AS "List<Struct<
            partner_id: String,
            token: String>>";

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

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

    DECLARE $token AS String;

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