from unittest.mock import MagicMock, patch

import pytest

from analytics.logic.demographics import get_demographics_2, get_table_for_demographics

MOCK_ROWS = [
    (
        11498668,
        72428294,
        113825769,
        160206201,
        133717727,
        46695882,
        7773334,
        1098867,
        274722368,
        252853614,
        19575030,
    )
]


RESULT = {
    "demographics": {
        "age": {
            "UA": 1098867,
            "_17": 11498668,
            "_18_22": 72428294,
            "_23_27": 113825769,
            "_28_34": 160206201,
            "_35_44": 133717727,
            "_45_59": 46695882,
            "_60_": 7773334,
        },
        "gender": {"F": 252853614, "M": 274722368, "UG": 19575030},
    },
    "account_id": 123,
    "account_type": "vendor",
    "sources": [
        {"id": 286, "name": "Spotify"},
        {"id": 1, "name": "Apple Music"},
        {"id": 716, "name": "Amazon Unlimited"},
    ],
}


@pytest.fixture
def mock_get_date_range():
    """Mock successful data_availability.get_date_range response."""
    with patch(
        "analytics.logic.demographics.data_availability.get_date_range"
    ) as get_date_range:
        get_date_range.return_value = ("2023-03-01", "2023-03-28")
        yield get_date_range


@pytest.fixture
def mock_add_outage_error_to_stores():
    """Mock add_outage_error_to_stores."""
    with patch(
        "analytics.logic.demographics.add_outage_error_to_stores"
    ) as add_outage_error_to_stores:
        add_outage_error_to_stores.return_value = [
            {"id": 286, "name": "Spotify"},
            {"id": 1, "name": "Apple Music"},
            {"id": 716, "name": "Amazon Unlimited"},
        ]
        yield add_outage_error_to_stores


@pytest.fixture
def mock_query_execute():
    with patch(
        "analytics.logic.demographics.AccountDemographics.execute",
        return_value=MOCK_ROWS,
    ) as query_execute:
        yield query_execute


@pytest.fixture
def mock_parallel(mock_add_outage_error_to_stores, mock_query_execute):
    result = MagicMock()
    result.message = {
        "demographics": mock_query_execute(),
        "sources": mock_add_outage_error_to_stores(),
    }
    with patch(
        "analytics.logic.demographics.parallel", return_value=result
    ) as mock_parallel:
        yield mock_parallel


def test_get_demographics_2(mock_parallel):
    result = get_demographics_2(
        {
            "account_id": 123,
            "query_type": "account_id",
            "account_type": "vendor",
            "start_date": "2020-01-01",
            "end_date": "2020-01-01",
            "is_apple_demographics_breakdown": False,
        },
        {
            "permission_subaccount_ids": [],
            "permission_label_ids": [],
            "permission_artist_ids": [],
            "permission_label_participant_ids": [],
        },
    )
    assert result == RESULT


@pytest.mark.parametrize(
    "query_type, countries, use_rollup_table, expected_table",
    [
        ("isrc", [], False, "V_STREAMS_DEMOGRAPHICS_BY_TRACK_FEED_DISTRIBUTOR_DAILY"),
        (
            "isrc",
            ["US"],
            False,
            "V_STREAMS_DEMOGRAPHICS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
        ),
        ("isrc", [], True, "STREAMS_DEMOGRAPHICS_BY_TRACK_FEED_DISTRIBUTOR_ROLLUP"),
        (
            "isrc",
            ["US"],
            True,
            "V_STREAMS_DEMOGRAPHICS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
        ),
        (
            "global_participant_id",
            [],
            False,
            "V_STREAMS_DEMOGRAPHICS_BY_PARTICIPANT_FEED_DISTRIBUTOR_DAILY",
        ),
        (
            "global_participant_id",
            ["US"],
            False,
            "V_STREAMS_DEMOGRAPHICS_BY_PARTICIPANT_COUNTRY_FEED_DISTRIBUTOR_DAILY",
        ),
        (
            "global_participant_id",
            [],
            True,
            "STREAMS_DEMOGRAPHICS_BY_PARTICIPANT_FEED_DISTRIBUTOR_ROLLUP",
        ),
        (
            "global_participant_id",
            ["US"],
            True,
            "STREAMS_DEMOGRAPHICS_BY_PARTICIPANT_COUNTRY_FEED_DISTRIBUTOR_ROLLUP",
        ),
        (
            "account_id",
            [],
            False,
            "V_STREAMS_DEMOGRAPHICS_BY_PRODUCT_FEED_DISTRIBUTOR_DAILY",
        ),
        (
            "account_id",
            ["US"],
            False,
            "V_STREAMS_DEMOGRAPHICS_BY_PRODUCT_COUNTRY_FEED_DISTRIBUTOR_DAILY",
        ),
        (
            "account_id",
            [],
            True,
            "STREAMS_DEMOGRAPHICS_BY_PRODUCT_FEED_DISTRIBUTOR_ROLLUP",
        ),
        (
            "account_id",
            ["US"],
            True,
            "STREAMS_DEMOGRAPHICS_BY_PRODUCT_COUNTRY_FEED_DISTRIBUTOR_ROLLUP",
        ),
    ],
)
def test_get_table_for_demographics(
    query_type, countries, use_rollup_table, expected_table
):
    """Test that get_table_for_demographics returns the correct table name."""
    table = get_table_for_demographics(query_type, countries, use_rollup_table)
    assert table == expected_table


def test_get_table_for_demographics_invalid_query_type():
    """Test that get_table_for_demographics raises ValueError for invalid query type."""
    with pytest.raises(ValueError, match="Invalid query type: invalid_type"):
        get_table_for_demographics("invalid_type", [], False)