# -*- coding: utf-8 -*-
from datacloud.dev_utils.yt import yt_config_table

PRETRAINED_MODELS_TABLE_PATH = '//home/x-products/production/config/datacloud/pretrained_models'


class PretrainedModelsTable(yt_config_table.ConfigTable):
    def __init__(self, table_path=PRETRAINED_MODELS_TABLE_PATH, yt_client=None):
        schema = [
            {'name': 'name', 'type': 'string', 'sort_order': 'ascending'},
            {'name': 'path', 'type': 'string'},
            {'name': 'features', 'type': 'string'},
            {'name': 'additional', 'type': 'any'},
        ]
        super(PretrainedModelsTable, self).__init__(
            table_path, schema, yt_client
        )

    def add_model(self, name, path, features, additional=None):
        additional = additional or {}

        record = {
            'name': name,
            'path': path,
            'features': features,
            'additional': additional,
        }
        self.insert_records([record])

    def get_model(self, name):
        return self.get_record_by_params({'name': name})
