from dataclasses import dataclass

import numpy as np
import pytest

from src.backend.constants import Flags, OutputColumns
from src.backend.constants import SnowFlakeColumns as SFCols
from src.backend.logic.audit.flags import base


class TestFlagChecks:
    _class = base._FlagChecks

    @pytest.mark.parametrize(
        "asset_id,should_flag",
        [
            ("", True),
            ("  ", True),
            (" abC", False),
            ("123", False),
            (None, True),
            (np.nan, True),
            (float("nan"), True),
        ],
    )
    def test_flag_no_asset_id(
        self, mock_flagger, mock_df, mocker_add_flag_db, asset_id, should_flag
    ):
        df = mock_df()
        df[SFCols.ASSET_ID] = asset_id
        instance = self._class(df, mock_flagger)
        instance.no_asset_id()
        was_flagged = df.loc[0][OutputColumns.AUDIT_FLAG_NO_ASSET_ID]

        assert was_flagged == should_flag
        if should_flag:
            mocker_add_flag_db_args = mocker_add_flag_db.call_args.args
            assert mocker_add_flag_db_args[0].empty != should_flag
            assert mocker_add_flag_db_args[1] == mock_flagger
            assert mocker_add_flag_db_args[2] == Flags.NO_ASSET_ID
        else:
            assert not mocker_add_flag_db.called


class TestFlagCleanup:
    _class = base._FlagCleanUp
    _flag_col = "flag"
    _sample_flags = [
        {
            "asset_id": "A123",
            "flag": "A",
        },
        {
            "asset_id": "A123",
            "flag": "B",
        },
        {
            "asset_id": "A123",
            "flag": "C",
        },
        {
            "asset_id": "A456",
            "flag": "A",
        },
        {
            "asset_id": "A456",
            "flag": "B",
        },
    ]

    @pytest.mark.parametrize("as_dataclasses", [True, False])
    def test_only_one_flag_flag_found(self, as_dataclasses):
        """Test that only_one_flag removes all flags except the one specified.
        This parameterized test checks as well that the behavior works both
        with dictionaries and dataclasses.
        """
        sample_flags = self._sample_flags.copy()
        flags = sample_flags + [{"asset_id": "A123", "flag": "X"}]

        if as_dataclasses:

            @dataclass
            class Flag:
                asset_id: str
                flag: str

            sample_flags = [Flag(**_) for _ in sample_flags]
            flags = [Flag(**_) for _ in flags]

        instance = self._class(flags, self._flag_col)
        assert instance.flags != sample_flags
        instance._only_one_flag("asset_id", flag="A", discard_flags=["X"])
        assert instance.flags == sample_flags

    def test_only_one_flag_flag_not_found(self):
        flags = self._sample_flags.copy() + [{"asset_id": "A123", "flag": "X"}]
        instance = self._class(flags, self._flag_col)
        assert instance.flags != self._sample_flags
        instance._only_one_flag("asset_id", flag="NOT_EXISTS", discard_flags=["X"])
        assert instance.flags == flags

    def test_only_one_flag_flags_to_discard_multiple(self):
        flags = self._sample_flags.copy() + [
            {"asset_id": "A123", "flag": f} for f in ["X", "X", "Y"]
        ]
        instance = self._class(flags, self._flag_col)
        assert instance.flags != self._sample_flags
        instance._only_one_flag("asset_id", flag="A", discard_flags=["X", "Y"])
        assert instance.flags == self._sample_flags

    def test_only_one_flag_flags_to_discard_empty(self):
        with pytest.raises(ValueError):
            self._class(self._sample_flags, self._flag_col)._only_one_flag(
                "asset_id", flag="A", discard_flags=[]
            )

    def test_only_one_flag_empty_df(self):
        instance = self._class([], self._flag_col)
        instance._only_one_flag("asset_id", flag="A", discard_flags=["X"])
        assert instance.flags == []