from enum import Enum
import yt.wrapper as yt_wrapper
from datacloud.config.yt import DATACLOUD_CONFIG_FOLDER
from datacloud.dev_utils.yt.yt_config_table import ConfigTable


BATCH_CONFIG_FOLDER = yt_wrapper.ypath_join(DATACLOUD_CONFIG_FOLDER, 'batch')
BATCH_TASK_DESCRIPTION_PATH = yt_wrapper.ypath_join(
    BATCH_CONFIG_FOLDER, 'description')


class StorageType(Enum):
    S3 = 's3'
    FTP = 'ftp'

    def __str__(self):
        return self.value


class TaskType(Enum):
    SCORING = 'scoring'

    def __str__(self):
        return self.value


class SubscriptionType(Enum):
    BATCH = 'batch'
    SIGN = 'sign'

    def __str__(self):
        return self.value


class BatchDescriptionTable(ConfigTable):
    def __init__(self, table_path=BATCH_TASK_DESCRIPTION_PATH, yt_client=None):
        schema = [
            {'name': 'task_id', 'type': 'string', 'sort_order': 'ascending'},
            {'name': 'partner_id', 'type': 'string'},
            {'name': 'task_type', 'type': 'string'},  # scoring/geo
            {'name': 'subscription_type', 'type': 'string'},  # batch/sign
            {'name': 'storage', 'type': 'string'},  # s3/ftp
            {'name': 'path', 'type': 'string'},
            {'name': 'data', 'type': 'any'},
            {'name': 'other', 'type': 'any'},
        ]
        super(BatchDescriptionTable, self).__init__(table_path, schema, yt_client)

    def add(self, task_id, partner_id, task_type, subscription_type,
            storage, path, data, other=None):
        other = other or {}
        assert isinstance(task_type, TaskType)
        assert isinstance(storage, StorageType)
        assert isinstance(subscription_type, SubscriptionType)
        record = {
            'task_id': task_id,
            'partner_id': partner_id,
            'task_type': str(task_type),
            'subscription_type': str(subscription_type),
            'storage': str(storage),
            'data': data,
            'path': path,
            'other': other,
        }
        self.insert_records([record])

    def get_description(self, task_id):
        return self.get_record_by_params({'task_id': task_id})
