from yt.wrapper import ypath_join, ypath_split
from datacloud.dev_utils.yt.yt_utils import create_folders
from datacloud.config.yt import MODELS_FOLDER
from datacloud.dev_utils.logging.logger import get_basic_logger
from datacloud.model_applyer.utils.helpers import (
    get_all_score_tables, BLENDS_FOLDER_NAME, BLENDS_TO_SAVE
)

logger = get_basic_logger(__name__)


def blend_reducer(_, recs):
    for rec in recs:
        yield rec
        break


def blend_scores(yt_client, partner_id, score_name, date_str):
    score_folder = ypath_join(MODELS_FOLDER, partner_id, score_name)
    blends_folder = ypath_join(score_folder, BLENDS_FOLDER_NAME)
    if not yt_client.exists(blends_folder):
        logger.info('Creating blends folder...')
        create_folders(blends_folder, yt_client=yt_client)

    blend_tables = sorted(yt_client.list(blends_folder, absolute=True))
    blends_count = len(blend_tables)
    assert blends_count <= BLENDS_TO_SAVE, 'Found too much blends in folder!'

    score_tables = get_all_score_tables(yt_client=yt_client, partner_id=partner_id,
                                        score_name=score_name)
    assert len(score_tables) > 0, 'No score tables found!'
    assert date_str == ypath_split(score_tables[-1])[-1], 'Bad date str!'

    if blends_count == 0:
        logger.info('No blends found. Initializing first blend...')
        tables_to_blend = score_tables[::-1]
    else:
        tables_to_blend = [score_tables[-1], blend_tables[-1]]

    logger.info('Tables to blend are {}'.format(tables_to_blend))

    table_schema = yt_client.get_attribute(score_tables[-1], 'schema')
    result_table = yt_client.TablePath(
        ypath_join(blends_folder, date_str),
        attributes={
            'schema': table_schema,
            'compression_codec': 'brotli_6',
            'erasure_codec': 'lrc_12_2_2',
            'optimize_for': 'scan',
        }
    )

    folder_prefix_len = len(score_folder) + 1
    if len(tables_to_blend) > 2:
        tables_to_blend_title = '{} ... {}'.foramt(
            tables_to_blend[0][folder_prefix_len:],
            tables_to_blend[-1][folder_prefix_len:]
        )
    else:
        tables_to_blend_title = ', '.join(
            t[folder_prefix_len:] for t in tables_to_blend
        )

    with yt_client.Transaction():
        yt_client.run_reduce(
            blend_reducer,
            tables_to_blend,
            result_table,
            reduce_by='cid',
            spec={'title': 'Blend score tables {}'.format(tables_to_blend_title)}
        )
        yt_client.run_sort(
            result_table,
            sort_by='cid',
            spec={'title': 'Blend score tables {} / sort after'.format(tables_to_blend_title)}
        )

        if len(blend_tables) == BLENDS_TO_SAVE:
            logger.info('Removing oldest blend {}'.format(blend_tables[0]))
            yt_client.remove(blend_tables[0])
