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


class PartnerScoresTable(YdbTable):
    def __init__(self, ydb_manager, database, table):
        super(PartnerScoresTable, 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('partner_score_name', ydb.OptionalType(ydb.DataType.String)))
                .with_column(ydb.Column('internal_score_name', ydb.OptionalType(ydb.DataType.String)))
                .with_column(ydb.Column('is_active', ydb.OptionalType(ydb.DataType.Bool)))
                .with_primary_keys('partner_id', 'partner_score_name')
            )

    def get_available_partner_scores(self, partner_id):
        query_params = {
            '$partner_id': partner_id
        }
        with self._init_session() as session:
            query = self._get_all_parther_scores_request.format(
                database=self.database,
                table=self.table
            )
            prepared_query = session.prepare(query)
            result_sets = session.transaction(ydb.SerializableReadWrite()).execute(
                prepared_query,
                query_params,
                commit_tx=True
            )
            for row in result_sets[0].rows:
                yield self.Record(**row)

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

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

    class Record(YdbTable.Record):
        __slots__ = ('partner_id', 'partner_score_name', 'internal_score_name', 'is_active')

        def __init__(self, partner_id, partner_score_name, internal_score_name=None, is_active=None):
            self.partner_id = partner_id
            self.partner_score_name = partner_score_name
            self.internal_score_name = internal_score_name
            self.is_active = is_active

        def __str__(self):
            return '({} {}; {}; {})'.format(self.partner_id, self.partner_score_name, self.internal_score_name, self.is_active)

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

        def __eq__(self, other):
            return self.partner_id == other.partner_id \
                and self.partner_score_name == other.partner_score_name \
                and self.internal_score_name == other.internal_score_name \
                and self.is_active == other.is_active

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

        DECLARE $records AS "List<Struct<
            partner_id: String,
            partner_score_name: String,
            internal_score_name: String,
            is_active: Bool>>";

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

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

    DECLARE $partner_id AS String;
    DECLARE $partner_score_name AS String;

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

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

    DECLARE $partner_id AS String;

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