import pytest

from delphi_api.core.constants import UTF8
from delphi_api.v3.data_models.bigtable.row_key_builder import RowKeyBuilder, RowKeyGroupType
from delphi_api.v3.enums import RowKeyAgg
from delphi_api.v3.view_models.params import Params
from tests.conftest import TEST_CHART_IDS, TEST_ISRCS, TEST_PLAYLIST_IDS
from tests.v3.constants import ARTIST_ID, END_DATE, START_DATE

_PLAYLIST_ID_1 = TEST_PLAYLIST_IDS[0]
_PLAYLIST_ID_2 = TEST_PLAYLIST_IDS[1]
_CHART_ID_1 = TEST_CHART_IDS[0]
_CHART_ID_2 = TEST_CHART_IDS[1]
_REV_ISRC_1 = TEST_ISRCS[0][::-1]
_REV_ISRC_2 = TEST_ISRCS[1][::-1]
_START_DATE = '2020-01-01'
_END_DATE = '2020-01-08'

TEST_ROW_KEY_PARAMS = [
    # @formatter:off
    (
        Params(isrc=TEST_ISRCS, start_date=START_DATE, end_date=END_DATE),
        RowKeyGroupType.STREAMS, RowKeyAgg.DAY,
        [(f'isrc_date_day~{_REV_ISRC_1}~{_START_DATE}', f'isrc_date_day~{_REV_ISRC_1}~{_END_DATE}'),
         (f'isrc_date_day~{_REV_ISRC_2}~{_START_DATE}', f'isrc_date_day~{_REV_ISRC_2}~{_END_DATE}')]
    ),
    (
        Params(isrc=TEST_ISRCS, playlist_id=TEST_PLAYLIST_IDS, start_date=START_DATE, end_date=END_DATE),
        RowKeyGroupType.STREAMS, RowKeyAgg.DAY,
        [(f'playlist_isrc_date_day~{_PLAYLIST_ID_1}~{_REV_ISRC_1}~{_START_DATE}', f'playlist_isrc_date_day~{_PLAYLIST_ID_1}~{_REV_ISRC_1}~{_END_DATE}'),
         (f'playlist_isrc_date_day~{_PLAYLIST_ID_1}~{_REV_ISRC_2}~{_START_DATE}', f'playlist_isrc_date_day~{_PLAYLIST_ID_1}~{_REV_ISRC_2}~{_END_DATE}'),
         (f'playlist_isrc_date_day~{_PLAYLIST_ID_2}~{_REV_ISRC_1}~{_START_DATE}', f'playlist_isrc_date_day~{_PLAYLIST_ID_2}~{_REV_ISRC_1}~{_END_DATE}'),
         (f'playlist_isrc_date_day~{_PLAYLIST_ID_2}~{_REV_ISRC_2}~{_START_DATE}', f'playlist_isrc_date_day~{_PLAYLIST_ID_2}~{_REV_ISRC_2}~{_END_DATE}')]
    ),
    (
        Params(isrc=TEST_ISRCS, playlist_id=TEST_PLAYLIST_IDS, start_date=START_DATE, end_date=END_DATE),
        RowKeyGroupType.PLAYLIST_POSITIONS, RowKeyAgg.DAY,
        [(f'playlist_isrc_date_day~{_PLAYLIST_ID_1}~{_REV_ISRC_1}~{_START_DATE}', f'playlist_isrc_date_day~{_PLAYLIST_ID_1}~{_REV_ISRC_1}~{_END_DATE}'),
         (f'playlist_isrc_date_day~{_PLAYLIST_ID_1}~{_REV_ISRC_2}~{_START_DATE}', f'playlist_isrc_date_day~{_PLAYLIST_ID_1}~{_REV_ISRC_2}~{_END_DATE}'),
         (f'playlist_isrc_date_day~{_PLAYLIST_ID_2}~{_REV_ISRC_1}~{_START_DATE}', f'playlist_isrc_date_day~{_PLAYLIST_ID_2}~{_REV_ISRC_1}~{_END_DATE}'),
         (f'playlist_isrc_date_day~{_PLAYLIST_ID_2}~{_REV_ISRC_2}~{_START_DATE}', f'playlist_isrc_date_day~{_PLAYLIST_ID_2}~{_REV_ISRC_2}~{_END_DATE}')]
    ),
    (
        Params(isrc=TEST_ISRCS, start_date=START_DATE, end_date=END_DATE),
        RowKeyGroupType.PLAYLIST_POSITIONS, RowKeyAgg.DAY,
        [(f'isrc_date_playlist_day~{_REV_ISRC_1}~{_START_DATE}', f'isrc_date_playlist_day~{_REV_ISRC_1}~{_END_DATE}'),
         (f'isrc_date_playlist_day~{_REV_ISRC_2}~{_START_DATE}', f'isrc_date_playlist_day~{_REV_ISRC_2}~{_END_DATE}')]
    ),
    (
        Params(isrc=TEST_ISRCS, chart_id=TEST_CHART_IDS, start_date=START_DATE, end_date=END_DATE),
        RowKeyGroupType.CHART_POSITIONS, RowKeyAgg.DAY,
        [(f'chart_isrc_date_day~{_CHART_ID_1}~{_REV_ISRC_1}~{_START_DATE}', f'chart_isrc_date_day~{_CHART_ID_1}~{_REV_ISRC_1}~{_END_DATE}'),
         (f'chart_isrc_date_day~{_CHART_ID_1}~{_REV_ISRC_2}~{_START_DATE}', f'chart_isrc_date_day~{_CHART_ID_1}~{_REV_ISRC_2}~{_END_DATE}'),
         (f'chart_isrc_date_day~{_CHART_ID_2}~{_REV_ISRC_1}~{_START_DATE}', f'chart_isrc_date_day~{_CHART_ID_2}~{_REV_ISRC_1}~{_END_DATE}'),
         (f'chart_isrc_date_day~{_CHART_ID_2}~{_REV_ISRC_2}~{_START_DATE}', f'chart_isrc_date_day~{_CHART_ID_2}~{_REV_ISRC_2}~{_END_DATE}')]
    ),
    (
        Params(isrc=TEST_ISRCS, dsp=['spotify'], start_date=START_DATE, end_date=END_DATE),
        RowKeyGroupType.CHART_POSITIONS, RowKeyAgg.DAY,
        [(f'dsp_isrc_date_chart_day~spotify~{_REV_ISRC_1}~{_START_DATE}', f'dsp_isrc_date_chart_day~spotify~{_REV_ISRC_1}~{_END_DATE}'),
         (f'dsp_isrc_date_chart_day~spotify~{_REV_ISRC_2}~{_START_DATE}', f'dsp_isrc_date_chart_day~spotify~{_REV_ISRC_2}~{_END_DATE}')]
    ),
    (
        Params(dsp=['spotify'], start_date=START_DATE, end_date=END_DATE),
        RowKeyGroupType.CHART_POSITIONS, RowKeyAgg.DAY,
        [(f'dsp_date_isrc_chart_day~spotify~{_START_DATE}', f'dsp_date_isrc_chart_day~spotify~{_END_DATE}')]
    ),
    (
        Params(isrc=TEST_ISRCS, start_date=START_DATE, end_date=END_DATE),
        RowKeyGroupType.CHART_POSITIONS, RowKeyAgg.DAY,
        [(f'isrc_date_chart_day~{_REV_ISRC_1}~{_START_DATE}', f'isrc_date_chart_day~{_REV_ISRC_1}~{_END_DATE}'),
         (f'isrc_date_chart_day~{_REV_ISRC_2}~{_START_DATE}', f'isrc_date_chart_day~{_REV_ISRC_2}~{_END_DATE}')]
    ),
    (
        Params(artist_id=ARTIST_ID, start_date=START_DATE, end_date=END_DATE),
        RowKeyGroupType.CHARTMETRIC, RowKeyAgg.DAILY,
        [(f'artist_daily~{ARTIST_ID}~{_START_DATE}', f'artist_daily~{ARTIST_ID}~{_END_DATE}')]
    ),
    # @formatter:on
]


class TestRowKeyBuilder:

    @pytest.mark.parametrize('params,key_group_type,agg,expected', TEST_ROW_KEY_PARAMS)
    def test_get_row_keys(self, params, key_group_type, agg, expected):
        result = RowKeyBuilder.get_row_keys(params, key_group_type, agg)
        encoded = []
        for keyset in expected:
            k1, k2 = keyset
            encoded.append((k1.encode(UTF8), k2.encode(UTF8),))
        assert result == encoded