import yt.wrapper as yt_wrapper
from datacloud.config.yt import UNRELIABLE_TMP_FOLDER as TMP
from datacloud.dev_utils.yt import yt_utils


def id_value_to_score(yt_client, input_table, score_table, snapshot, output_table):
    assert yt_utils.check_table_exists(score_table, yt_client)

    if not yt_utils.is_sorted_by(yt_client, input_table, ['id_value']):
        yt_client.run_sort(input_table, sort_by='id_value')

    if not yt_client.exists(TMP):
        yt_utils.create_folders([TMP], yt_client)
    with yt_client.TempTable(TMP) as id_to_cid_tmp_table, \
         yt_client.TempTable(TMP) as id_to_score_tmp_table:
        _match_id_val_to_cid(yt_client, input_table, snapshot, id_to_cid_tmp_table)
        yt_client.run_sort(id_to_cid_tmp_table, sort_by='cid')
        _match_to_score(yt_client, id_to_cid_tmp_table, score_table,
                        id_to_score_tmp_table)
        _set_unique_score(yt_client, id_to_score_tmp_table, output_table)


@yt_wrapper.with_context
def _id_value_to_cid_reducer(_, recs, context):
    ext_id = None
    for rec in recs:
        if context.table_index == 0:
            ext_id = rec['external_id']
        elif ext_id:
            yield {
                'external_id': ext_id,
                'id_value': rec['id_value'],
                'cid': rec['cid'],
            }
        else:
            break


def _match_id_val_to_cid(yt_client, input_table, crypta_snapshot, output_table):
    yt_client.run_reduce(
        _id_value_to_cid_reducer,
        [
            input_table,
            crypta_snapshot.email_id_value_to_cid,
            crypta_snapshot.phone_id_value_to_cid,
        ],
        output_table,
        reduce_by='id_value',
    )


@yt_wrapper.with_context
def _to_score_reducer(_, recs, context):
    score = None
    for rec in recs:
        if context.table_index == 0:
            score = rec['score']
        elif score is not None:
            yield {
                'external_id': rec['external_id'],
                'score': score,
            }
        else:
            break


def get_score_table_from(score_info, date_str):
    raise NotImplementedError('TODO: Implement')


def _match_to_score(yt_client, input_table, score_table, output_table):
    yt_client.run_reduce(
        _to_score_reducer,
        [score_table, input_table],
        output_table,
        reduce_by='cid',
    )


def _select_max_reducer(_, recs):
    score = None
    for rec in recs:
        if score is None:
            score, ext_id = rec['score'], rec['external_id']
        else:
            score = max(score, rec['score'])
    yield {'external_id': ext_id, 'score': score}


def _set_unique_score(yt_client, input_table, output_table):
    yt_client.run_map_reduce(
        None,
        _select_max_reducer,
        input_table,
        output_table,
        reduce_by='external_id'
    )
