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


class RejectedRequestsTable(YdbTable):
    def __init__(self, ydb_manager, database, table):
        super(RejectedRequestsTable, 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('partner_id', ydb.OptionalType(ydb.DataType.String)))
                .with_column(ydb.Column('score_id', ydb.OptionalType(ydb.DataType.String)))
                .with_column(ydb.Column('request_id', ydb.OptionalType(ydb.DataType.String)))
                .with_column(ydb.Column('request_timestamp', ydb.OptionalType(ydb.DataType.Uint64)))
                .with_column(ydb.Column('reject_timestamp', ydb.OptionalType(ydb.DataType.Uint64)))
                .with_primary_key('partner_id')
                .with_primary_key('score_id')
                .with_primary_key('request_id')
            )

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

    def get(self, record):
        raise '`get` is not supported by RejectTable, use `get_one`.'

    def get_multiple(self, _):
        raise '`get_multiple` is not supported by RejectTable, use `get_one`.'

    class Record(object):
        __slots__ = ('partner_id', 'score_id', 'request_id', 'request_timestamp', 'reject_timestamp')

        def __init__(self, partner_id, score_id, request_id, request_timestamp=None, reject_timestamp=None):
            self.partner_id = partner_id
            self.score_id = score_id
            self.request_id = request_id
            self.request_timestamp = request_timestamp
            self.reject_timestamp = reject_timestamp

        def __str__(self):
            return '(({partner_id}, {score_id}, {request_id}): {request_timestamp}, {reject_timestamp})'.format(
                partner_id=self.partner_id,
                score_id=self.score_id,
                request_id=self.request_id,
                request_timestamp=self.request_timestamp,
                reject_timestamp=self.reject_timestamp
            )

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

        def __eq__(self, other):
            return self.partner_id == other.partner_id and \
                self.score_id == other.score_id and \
                self.request_id == other.request_id and \
                self.request_timestamp == other.request_timestamp and \
                self.reject_timestamp == other.reject_timestamp

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

        DECLARE $records AS "List<Struct<
            partner_id: String,
            score_id: String,
            request_id: String,
            request_timestamp: Uint64,
            reject_timestamp: Uint64>>";

        REPLACE INTO [{table}]
        SELECT
            partner_id,
            score_id,
            request_id,
            request_timestamp,
            reject_timestamp
        FROM AS_TABLE($records);
    """

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

        DECLARE $partner_id AS String;
        DECLARE $score_id AS String;
        DECLARE $request_id AS String;

        SELECT *
        FROM [{table}]
        WHERE partner_id = $partner_id AND score_id = $score_id AND request_id = $request_id;
    """
