# -*- coding: utf-8 -*-
import json
from datacloud.dev_utils.logging.logger import get_basic_logger
from datacloud.dev_utils.yt import yt_utils
from datacloud.dev_utils.yt.converter import TableConverter
from datacloud.dev_utils.time import patterns, utils as time_utils
from datacloud.dev_utils.transfer.yt_to_ydb import table_description
from datacloud.dev_utils.transfer.yt_to_ydb import transfer_table
from datacloud.dev_utils.yql import yql_helpers
from datacloud.config.yt import YT_PROXY

logger = get_basic_logger(__name__)


class YdbStabilityTableDescription(table_description.YdbTableDescription):
    def __init__(self, table_path, ydb_connection_params):
        super(YdbStabilityTableDescription, self).__init__(table_path, ydb_connection_params)
        self.schema = """
        (
            date String,
            partner_id String,
            score_id String,
            segment String,
            stability String,
            primary key (date, partner_id, score_id, segment)
        )
        """
        self.fields = """
            `date`,
            partner_id,
            score_id,
            segment,
            stability
        """


class CStabilityToTransferRec(TableConverter):
    def __init__(self, yt_client, partner_id, score_id, date_str):
        super(CStabilityToTransferRec, self).__init__(yt_client)
        time_utils.assert_date_str(date_str, patterns.RE_DAILY_LOG_FORMAT)

        self._date = date_str
        self._partner_id = partner_id
        self._score_id = score_id
        self._output_table_schema = [
            {'name': 'partner_id', 'type': 'string'},
            {'name': 'score_id', 'type': 'string'},
            {'name': 'date', 'type': 'string'},
            {'name': 'segment', 'type': 'string'},
            {'name': 'stability', 'type': 'string'}
        ]

    def _convert(self, input_table, output_table):
        output_table = yt_utils.TablePath(output_table, schema=self._output_table_schema)

        with self._yt_client.Transaction():
            segments = self._get_unique_segments(input_table)
            records_to_upload = []
            with self._yt_client.TempTable() as sorted_input_table:
                self._yt_client.run_sort(
                    input_table,
                    sorted_input_table,
                    sort_by=('segment')
                )
                for segment in segments:
                    records_to_upload.append(self._prepare_stability_record(sorted_input_table, segment))
            self._yt_client.write_table(output_table, records_to_upload)

    def _prepare_stability_record(self, table, segment):
        data = {}
        for row in self._yt_client.read_table(yt_utils.TablePath(table, exact_key=(segment, ))):
            data[row['feature']] = row['hist']
        data = [data[key] for key in sorted(data)]
        data = json.dumps(data)
        return {'date': self._date, 'partner_id': self._partner_id, 'score_id': self._score_id, 'segment': segment, 'stability': data}

    def _get_unique_segments(self, segment_table):
        available_segments = set()
        for row in self._yt_client.read_table(yt_utils.TablePath(segment_table, columns=['segment'])):
            available_segments.add(row['segment'])
        return available_segments


def sample_stability_trasnfer(yt_client, yt_table, ydb_table, yql_token=None):
    yql_client = yql_helpers.create_yql_client(yt_utils.get_yt_client(), token=yql_token)

    ydb_connection_params = table_description.YdbConnectionParams(
        endpoint='ydb-ru.yandex.net:2135',
        database='/ru/impulse/production/xprod-scoring-api-db'
    )

    source_table = table_description.YtTableDescription(
        table_path=yt_table,
        yt_cluster=YT_PROXY
    )
    target_table = YdbStabilityTableDescription(
        table_path=ydb_table,
        ydb_connection_params=ydb_connection_params
    )

    transfer_table.transfer(yql_client, source_table, target_table, yql_token)


if __name__ == '__main__':
    _yt_client = yt_utils.get_yt_client()

    date_str = '2019-01-23'
    partner_id = 'tcs'
    score_id = 'tcs_xprod_987_20181212'

    import os  # noqa
    _yql_token = os.environ['YQL_TOKEN']

    CStabilityToTransferRec(_yt_client, partner_id, score_id, date_str, 'daily')(
        '//home/x-products/production/datacloud/stability/tcs/tcs_xprod_987_20181212/daily/{date_str}'.format(date_str=date_str),
        '//tmp/r9-stability-rec'
    )

    sample_stability_trasnfer(
        _yt_client,
        '//tmp/r9-stability-rec',
        '/ru/impulse/production/xprod-scoring-api-db/score_api/stability/stability',  # 'tcs/score_name/stability',
        _yql_token
    )
