"""Flag checks for AT audits."""

import asyncio
import re
from functools import lru_cache
from typing import Iterable, Self

import pandas as pd

from ....constants import Countries, Flags, YTMisc
from ....logic.audit.logic.common import OutputColumns, SFCols
from ....typings import ArtistName, ChannelName, CountryCode2, RowIndex
from ....utils.html import bold
from ....utils.strings import remove_accents, remove_multispace
from ..logic import common
from . import conditions
from .base import _FlagChecks


class FlagChecksAT(_FlagChecks):
    # Ownership will be considered valid if those territories are the only ones missing.
    # This is a special case for Art Tracks and has been explicitly requested
    # by the Ops / Audit team.
    VALID_OWNERSHIP_EXTRA_TERRITORIES: set[CountryCode2] = {
        "BL",
        "BQ",
        "CW",
        "MF",
        "SS",
        "SX",
    }

    def no_topic_channel(self) -> Self:
        """Flag in place rows without topic channel. Rows without Asset ID
        will also be flagged by this check.
        """
        at = self.df
        flag_col = OutputColumns.AUDIT_FLAG_NO_TOPIC_CHANNEL
        no_topic_channel = conditions.has_topic_channel(at).eq(False)
        no_asset_id = conditions.has_asset_id(at).eq(False)

        at[flag_col] = no_topic_channel | no_asset_id

        flag_rows = at[at[flag_col]]
        if len(flag_rows):
            common.add_flag_db(
                flag_rows,
                self.flagger,
                Flags.NO_TOPIC_CHANNEL,
            )
        return self

    def bad_topic_channel(self) -> Self:
        """Flag in place rows with bad topic channel.

        The channel name should be included in the artist name string (case
        insensitive), with the suffix "Topic" removed, or vice versa.

        Example: "Artist - Topic" -> "Artist, Producer".

        Does not perform exact full string matching, as the channel name may be
        slightly different from the artist name (e.g. "Artist" vs "Artist,
        Producer"). Most common accents are removed from the strings before
        comparison, to make the comparison accent-insensitive.

        Audits team docs on which this implementation is based:
        https://docs.google.com/document/d/1uQDDb6zuURuhBIDZvPEyNsiBo7sIjwlK259lhaOhdLE
        """

        @lru_cache(maxsize=1024)
        def syntactic_check(channel_name: ChannelName, artist_name: ArtistName) -> bool:
            """Check if the artist name is included in the channel name string,
            or vice versa, with the "Topic" suffix removed from the channel name and
            in a case insensitive manner.

            Accent discrepancies are ignored on purpose.

            Args:
                channel_name (str): Channel name.
                artist_name (str): Artist name.

            Returns:
                bool: True if the channel name is included in the artist name,
                    or vice versa, False otherwise.
            """

            def preprocess(string: str) -> str:
                unaccented_lower = remove_accents(string).lower()
                ex_punctuation = re.sub(r"[^a-z0-9\s]", "", unaccented_lower)
                no_multiple_spaces_stripped = remove_multispace(ex_punctuation).strip()
                return no_multiple_spaces_stripped

            # Preprocess the channel name and artist name.
            preprocessed_channel_name = preprocess(
                str(channel_name)
                .lower()
                .replace(YTMisc.TOPIC_CHANNEL_SUFFIX.lower(), "")
            )
            preprocessed_artist_name = preprocess(artist_name)
            return (preprocessed_artist_name in preprocessed_channel_name) or (
                preprocessed_channel_name in preprocessed_artist_name
            )

        at = self.df
        flag_col = OutputColumns.AUDIT_FLAG_BAD_TOPIC_CHANNEL
        at[flag_col] = at.apply(
            # No flag if either channel name or artist name is missing.
            # If channel name is missing, the flag will be set in the NO_TOPIC_CHANNEL check.
            lambda row: isinstance(row[SFCols.CHANNEL_DISPLAY_NAME], str)
            and isinstance(row[SFCols.ARTIST], str)
            and not syntactic_check(
                row[SFCols.CHANNEL_DISPLAY_NAME], row[SFCols.ARTIST]
            ),
            axis=1,
        )
        flag_rows = at[at[flag_col]]
        if len(flag_rows):
            common.add_flag_db(
                flag_rows,
                self.flagger,
                Flags.BAD_TOPIC_CHANNEL,
                details=lambda x: (
                    f"Current is "
                    f"{bold(x[SFCols.CHANNEL_DISPLAY_NAME].removesuffix(YTMisc.TOPIC_CHANNEL_SUFFIX))}, "
                    f"should be {bold(x[SFCols.ARTIST])}",
                ),
            )
        return self

    # Check the special valid ownership type for Art Tracks.
    # There shouldn't be any territories other than the ones in the set
    # VALID_OWNERSHIP_EXTRA_TERRITORIES or RU missing (the "ignored territories").
    ignore_territories = {Countries.RU} | VALID_OWNERSHIP_EXTRA_TERRITORIES

    async def ownership_incomplete(self) -> None:
        """Flag in place rows with incomplete ownership / ownership not
        in the list of valid ownership types.
        """
        at = self.df

        instance = _OwnershipIncomplete(at, self.ignore_territories)
        await instance.run()

        flag_rows = at[at[OutputColumns.AUDIT_FLAG_OWNERSHIP_INCOMPLETE].eq(True)]

        def _ownership_incomplete_details(row_idx: RowIndex) -> str | None:
            missing_territories = instance.missing_territories.get(row_idx)
            if missing_territories:
                return f"Missing territories: {bold(', '.join(sorted(missing_territories)))}"
            return None

        if len(flag_rows):
            common.add_flag_db(
                flag_rows,
                self.flagger,
                Flags.OWNERSHIP_INCOMPLETE,
                details=lambda row: _ownership_incomplete_details(row.name),
            )


class _OwnershipIncomplete:
    """Class to encapsulate the ownership incomplete flag check.

    Attributes:

        _ignore_territories_in_cols: Columns with 2-letter country codes, countries
            which should be ignored in the ownership incomplete check.

        _flag_col: Flag column for the ownership incomplete check. This column will be
            added to the DataFrame if it doesn't exist and will be set to True for rows
            with incomplete ownership.

        at: DataFrame with the AT data.

        ignore_territories: Territories to ignore for all rows in the AT DataFrame.

        missing_territories: Dictionary with the missing territories for each row.
            They are set upon running the check and are used in the flag details.
            The keys are the row indices and the values are lists of missing
            territories. Rows with no missing territories will not be included in
            this dictionary for memory efficiency.
    """

    _ignore_territories_in_cols: set[SFCols] = {
        SFCols.ABBRIVATION,  # Carved out
        SFCols.VTR_COUNTRIES,  # Vendor Territory Restrictions
        SFCols.RTR_COUNTRIES,  # Release Territory Restrictions
        SFCols.STR_COUNTRIES,  # Subaccount Territory Restrictions
    }

    _flag_col = OutputColumns.AUDIT_FLAG_OWNERSHIP_INCOMPLETE

    def __init__(
        self, at: pd.DataFrame, ignore_territories: Iterable[CountryCode2] | None = None
    ):
        """
        Initialize the ownership incomplete flag check. All changes are made in place
        in the provided `at` DataFrame.

        Args:
            at: DataFrame with the AT data.
            ignore_territories: Territories to ignore for all rows in the DataFrame
                (i.e. ignore them for all labels, releases, etc.).
                Those are territories that if missing from the ownership type
                will not trigger the flag. Default is None.
        """
        self.at = at
        self.ignore_territories = frozenset(ignore_territories or [])
        self.missing_territories: dict[RowIndex, list[CountryCode2]] = {}

        # Add the flag column if it doesn't exist
        if self._flag_col not in at.columns:
            at[self._flag_col] = pd.NA

    @property
    def rows_with_asset_id(self) -> pd.DataFrame:
        """Get the rows with asset ID (as a view)."""
        has_asset_id = conditions.has_asset_id(self.at)
        return self.at[has_asset_id]

    async def run(self) -> None:
        """Flag rows with incomplete ownership or ownership not in valid types. Do
        not flag rows with missing asset ID, as those will be flagged by other
        checks.

        Returns:
            None, as the flag is added in place.
        """
        # For I/O performance reasons, use a row-based approach with asyncio.
        await asyncio.gather(*map(self._handle_row, self.rows_with_asset_id.index))

    async def _handle_row(self, idx: RowIndex) -> None:
        """Check if ownership is valid and complete. Ownership is valid if it
        is a set and is complete if it includes all valid territories,
        except for certain to be ignored.

        Args:
            idx: Index of the row to check.

        Returns:
            None, as the flag is added in place.
        """
        row = self.at.loc[idx]
        row_ownership_territories = row[SFCols.OWNERSHIP]

        # Ownership is not valid if it's not a set (e.g. NaN, None).
        # Flag the row in this case, and short-circuit the function.
        if not isinstance(row_ownership_territories, set):
            self._flag_row(idx)
            return

        row_territories_to_ignore = self._get_row_territories_to_ignore(row)

        ownership_valid_check = await common.is_ownership_valid(
            row_ownership_territories,
            ignore_territories=row_territories_to_ignore,
        )
        is_ownership_valid, missing_territories = ownership_valid_check

        if not is_ownership_valid:
            # Persist the missing territories for the row, to use them
            # in the flag details.
            self.missing_territories[idx] = missing_territories
            self._flag_row(idx)

    def _flag_row(self, idx: RowIndex) -> None:
        """Flag a row, setting the flag column to True."""
        self.at.at[idx, self._flag_col] = True

    def _get_row_territories_to_ignore(self, row: pd.Series) -> set[CountryCode2]:
        """Get the territories to ignore for a row. Territories to ignore are
        those provided in the constructor and those in the columns defined in
        `_ignore_territories_in_cols`.

        Args:
            row: Row of the DataFrame to get the territories to ignore for.
        """
        territories_to_ignore = set() | self.ignore_territories
        for col in self._ignore_territories_in_cols:
            value = row[col]
            if pd.notna(value) and value:
                territories_to_ignore.update(value)  # Assume value is a set.
        return territories_to_ignore