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


class StabilityTable(YdbTable):
    def __init__(self, ydb_manager, database, table):
        # TODO: change score data to correct balue "$scoreData"
        super(StabilityTable, self).__init__(ydb_manager, database, table)

    def create(self):
        with self._init_session() as session:
            # TODO: Never was tested
            session.create_table(
                self.full_table_path,
                ydb.TableDescription()
                .with_column(ydb.Column('date', ydb.OptionalType(ydb.DataType.String)))
                .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('segment', ydb.OptionalType(ydb.DataType.String)))
                .with_column(ydb.Column('stability', ydb.OptionalType(ydb.DataType.String)))
                .with_primary_key('date')
                .with_primary_key('partner_id')
                .with_primary_key('score_id')
                .with_primary_key('segment')
            )

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

    # TODO: Never was tested
    def get(self, record):
        query_params = {
            '$date': record.date,
            '$partner_id': record.partner_id,
            '$score_id': record.score_id,
            '$segment': record.segment
        }
        for rec in self._get(query_params):
            yield rec

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

    class Record(object):
        __slots__ = ('date', 'partner_id', 'score_id', 'segment', 'stability')

        def __init__(self, date, partner_id, score_id, segment, stability=None):
            self.date = date
            self.partner_id = partner_id
            self.score_id = score_id
            self.segment = segment
            self.stability = stability

        def __str__(self):
            return '(({date}, {partner_id}, {score_id}, {segment}): {stability})'.format(
                date=self.date,
                partner_id=self.partner_id,
                score_id=self.score_id,
                segment=self.segment,
                stability=self.stability
            )

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

        def __eq__(self, other):
            return self.date == other.date and \
                self.partner_id == other.partner_id and \
                self.score_id == other.score_id and \
                self.segment == other.segment and \
                self.stability == other.stability

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

        DECLARE $records AS "List<Struct<
            date: String,
            partner_id: String,
            score_id: String,
            segment: String,
            stability: String>>";

        REPLACE INTO [{table}]
        SELECT
            date,
            partner_id,
            score_id,
            segment,
            stability
        FROM AS_TABLE($records);
    """

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

        DECLARE $date AS String;
        DECLARE $partner_id AS String;
        DECLARE $score_id AS String;
        DECLARE $segment AS String;

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