# -*- coding: utf-8 -*-
from yt.wrapper import ypath_join
from datacloud.config.yt import PRODUCTION_ROOT
from datacloud.dev_utils.yt.yt_config_table import ConfigTable


DEFAULT_TIME_DB_PATH = ypath_join(PRODUCTION_ROOT, 'metrics/task-time-log')


class TimeRecord:
    __slots__ = ('program', 'key', 'status', 'duration', 'other')

    def __init__(self, program, timestamp, key, status, duration, other=None):
        self.program = program
        self.timestamp = int(timestamp)
        self.key = key
        self.status = status
        self.duration = duration
        self.other = other or {}

    def to_json(self):
        return {
            'program': self.program,
            'timestamp': int(self.timestamp),
            'key': self.key,
            'status': str(self.status),
            'duration': self.duration,
            'other': self.other
        }


class TimeDB(ConfigTable):
    def __init__(self, table_path=DEFAULT_TIME_DB_PATH, yt_client=None):
        schema = [
            {'name': 'program', 'type': 'string', 'sort_order': 'ascending', 'required': True},
            {'name': 'timestamp', 'type': 'uint64', 'sort_order': 'ascending', 'required': True},
            {'name': 'key', 'type': 'string', 'sort_order': 'ascending', 'required': True},
            {'name': 'status', 'type': 'string', 'required': True},
            {'name': 'duration', 'type': 'uint64', 'required': True},
            {'name': 'other', 'type': 'any', 'required': False}
        ]

        super(TimeDB, self).__init__(
            table_path, schema, yt_client
        )

    def add(self, time_record):
        self.add_multiple([time_record])

    def add_multiple(self, time_records):
        records = [rec.to_json() for rec in time_records]
        self.insert_records_with_retry(records, n_retry=10)


if __name__ == '__main__':
    from datacloud.dev_utils.yt import yt_utils
    TimeDB(yt_client=yt_utils.get_yt_client()).create_table()
