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


class GeoPathTable(YdbTable):
    def __init__(self, ydb_manager, database, table):
        super(GeoPathTable, 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('internal_score_name', ydb.OptionalType(ydb.DataType.String)))
                .with_column(ydb.Column('crypta_path', ydb.OptionalType(ydb.DataType.String)))
                .with_column(ydb.Column('geo_path', ydb.OptionalType(ydb.DataType.String)))
                .with_column(ydb.Column('weights', ydb.OptionalType(ydb.DataType.String)))
                .with_primary_key('internal_score_name')
            )

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

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

    class Record(object):
        __slots__ = ('internal_score_name', 'crypta_path', 'geo_path', 'weights')

        def __init__(self, internal_score_name, crypta_path=None, geo_path=None, weights=None):
            self.internal_score_name = internal_score_name
            self.crypta_path = crypta_path
            self.geo_path = geo_path
            self.weights = weights

        def __str__(self):
            return '({}: {}, {}, {})'.format(
                self.internal_score_name,
                self.crypta_path,
                self.geo_path,
                self.weights
            )

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

        def __eq__(self, other):
            return self.internal_score_name == other.internal_score_name and self.geo_path == other.geo_path \
                and self.crypta_path == other.crypta_path and self.weights == other.weights

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

        DECLARE $records AS "List<Struct<
        internal_score_name: String,
        crypta_path: String,
        geo_path: String,
        weights: String>>";

        REPLACE INTO [{table}]
        SELECT
            internal_score_name,
            geo_path,
            crypta_path,
            weights
        FROM AS_TABLE($records);
    """

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

        DECLARE $internal_score_name AS String;

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