# -*- coding: utf-8 -*-
import murmurhash
import json
import datetime
from collections import namedtuple

from datacloud.dev_utils.yt import yt_config_table
from datacloud.dev_utils.time.patterns import FMT_DATE_HMST


ExpsStorageRec = namedtuple('ExpsStorageRec', [
    'key',
    'config',
    'metrics',
    'timestamp',
    'additional',
])


class BaseExpsStorage(object):
    def _make_exp_key(self, config):
        raise NotImplementedError()

    def get_key(self, config):
        return self._make_exp_key(config)

    def _write(self, key, config, metrics, additional=None, force=False):
        raise NotImplementedError()

    def add_experiment(self, config, metrics, force=False):
        key = self._make_exp_key(config=config)
        rec = self.read(key=key)
        if force or rec is None:
            return self._write(key=key, config=config, metrics=metrics, additional={})
        return rec

    def _read(self, key):
        raise NotImplementedError()

    def read(self, key):
        if key is None:
            return None
        return self._read(key=key)

    def read_by_config(self, config):
        key = self._make_exp_key(config=config)
        return self.read(key=key)


def serialize_config(config):
    return json.dumps(config, sort_keys=True)


def hash_config(serialized_config):
    return murmurhash.hash64(serialized_config)


YT_EXPS_STORAGE_PATH = '//home/x-products/penguin-dvier-dev/cache/exps_storage'


class YtExpsStorage(BaseExpsStorage, yt_config_table.ConfigTable):
    class YtTable(yt_config_table.ConfigTable):
        def __init__(self, table_path, yt_client):
            schema = [
                {'name': 'key', 'type': 'uint64', 'sort_order': 'ascending'},
                {'name': 'config', 'type': 'any'},
                {'name': 'metrics', 'type': 'any'},
                {'name': 'timestamp', 'type': 'string'},
                {'name': 'additional', 'type': 'any'},
            ]
            super(YtExpsStorage.YtTable, self).__init__(
                table_path=table_path, schema=schema, yt_client=yt_client
            )

    def __init__(self, table_path=YT_EXPS_STORAGE_PATH, yt_client=None):
        self._yt_table = YtExpsStorage.YtTable(table_path=table_path, yt_client=yt_client)

    def _make_exp_key(self, config):
        serialized_config = serialize_config(config)
        return hash_config(serialized_config)

    def _write(self, key, config, metrics, additional=None):
        additional = additional or {}
        record = {
            'key': key,
            'config': config,
            'metrics': metrics,
            'timestamp': datetime.datetime.now().strftime(FMT_DATE_HMST),
            'additional': additional
        }
        self._yt_table.insert_records([record])
        return ExpsStorageRec(**record)

    def _read(self, key):
        record = self._yt_table.get_record_by_params({'key': key})
        if record is None:
            return None
        return ExpsStorageRec(**record)


class BaseBenchmarksStorage(object):
    def _make_bench_key(self, configs):
        raise NotImplementedError()

    def _write(self, key, configs, metrics, additional=None):
        raise NotImplementedError()

    def _read(self, key):
        raise NotImplementedError()

    def read(self, key):
        if key is None:
            return None
        return self._read(key=key)

    def add_benchmark(self, configs, metrics, force=False):
        key = self._make_bench_key(configs=configs)
        if force or self.read(key=key) is None:
            self._write(key=key, configs=configs, metrics=metrics, additional={})

    def read_by_configs(self, configs):
        key = self._make_bench_key(configs=configs)
        return self.read(key=key)


class YtBenchmarksStorage(BaseExpsStorage, yt_config_table.ConfigTable):
    class YtTable(yt_config_table.ConfigTable):
        def __init__(self, table_path, yt_client):
            schema = [
                {'name': 'key', 'type': 'string', 'sort_order': 'ascending'},
                {'name': 'configs', 'type': 'any'},
                {'name': 'metrics', 'type': 'any'},
                {'name': 'timestamp', 'type': 'string'},
                {'name': 'additional', 'type': 'any'},
            ]
            super(YtBenchmarksStorage.YtTable, self).__init__(
                table_path=table_path, schema=schema, yt_client=yt_client
            )

    def __init__(self, table_path, yt_client=None, separator='-'):
        self.separator = separator
        self._yt_table = YtBenchmarksStorage.YtTable(table_path=table_path, yt_client=yt_client)

    def _make_bench_key(self, configs):
        serialized_configs = [serialize_config(c) for c in configs]
        configs_hashes = [hash_config(sc) for sc in serialized_configs]
        assert len(set(configs_hashes)) == len(configs), 'All configs should be different'

        sorted_configs = sorted(configs_hashes)
        return self.separator.join(map(str, sorted_configs))

    def _write(self, key, configs, metrics, additional=None):
        additional = additional or {}
        record = {
            'key': key,
            'configs': configs,
            'metrics': metrics,
            'timestamp': datetime.datetime.now().strftime(FMT_DATE_HMST),
            'additional': additional
        }
        self._yt_table.insert_records([record])

    def _read(self, key):
        return self._yt_table.get_record_by_params({'key': key})
