from enum import Enum
import yt.wrapper as yt_wrapper
from datacloud.config.yt import DATACLOUD_CONFIG_FOLDER
from datacloud.dev_utils.time.patterns import FMT_DATE_HMS, RE_FMT_DATE_HMS
from datacloud.dev_utils.time import utils as time_utils
from datacloud.dev_utils.yt.yt_config_table import ConfigTable


BATCH_CONFIG_FOLDER = yt_wrapper.ypath_join(DATACLOUD_CONFIG_FOLDER, 'batch')
BATCH_LOG_TABLE_PATH = yt_wrapper.ypath_join(BATCH_CONFIG_FOLDER, 'log')


class Status(Enum):
    NEW = 'new'
    LOST = 'lost'
    INVALID = 'invalid'
    # ACCEPTED = 'accepted'
    SCHEDULED = 'scheduled'
    WAIT_UPLOAD = 'wait_upload'
    DONE = 'done'

    def __str__(self):
        return self.value


class BatchStatusTable(ConfigTable):
    def __init__(self, table_path=BATCH_LOG_TABLE_PATH, yt_client=None):
        schema = [
            {'name': 'task_id', 'type': 'string', 'sort_order': 'ascending'},
            {'name': 'key', 'type': 'string', 'sort_order': 'ascending'},
            {'name': 'status', 'type': 'string'},
            {'name': 'add_time', 'type': 'string'},
            {'name': 'update_time', 'type': 'string'},
            {'name': 'data', 'type': 'any'},
            {'name': 'other', 'type': 'any'},
        ]
        super(BatchStatusTable, self).__init__(
            table_path, schema, yt_client)

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

    def get_existing_input_names(self, task_id):
        items = set()
        for rec in self.get_records_by_params({'task_id': task_id}):
            items.add(rec.get('data', {}).get('input_name'))
        return items

    def get_existing_keys(self, task_id):
        items = set()
        for rec in self.get_records_by_params({'task_id': task_id}):
            items.add(rec.get('key'))
        return items

    def get_with_status(self, task_id, status):
        assert isinstance(status, Status)
        for item in self.get_records_by_params({
                'task_id': task_id, 'status': str(status)}):
            yield item

    def update_status(self, task_id, key, new_status, update_time=None):
        assert isinstance(new_status, Status)
        new_status = str(new_status)
        update_time = update_time or time_utils.now_str(FMT_DATE_HMS)
        time_utils.assert_date_str(update_time, pattern=RE_FMT_DATE_HMS)

        item = self.get_record_by_params({'task_id': task_id, 'key': key})
        item['status'] = new_status
        item['update_time'] = update_time
        self.insert_records([item])

    def add(self, task_id, key, data,
            status=Status.NEW, add_time=None,
            update_time=None, other=None):
        assert isinstance(status, Status)
        current_time = time_utils.now_str(FMT_DATE_HMS)
        add_time = add_time or current_time
        update_time = update_time or current_time
        time_utils.assert_date_str(add_time, pattern=RE_FMT_DATE_HMS)
        time_utils.assert_date_str(update_time, pattern=RE_FMT_DATE_HMS)
        other = other or {}
        record = {
            'task_id': task_id,
            'key': key,
            'data': data,
            'status': str(status),
            'add_time': add_time,
            'update_time': update_time,
            'other': other,
        }
        self.insert_records([record])
