import asyncio
import base64
from asyncio import Task
from io import BytesIO
from typing import Callable

import pandas as pd

from .... import logger
from ....connectors.db.db import Client as DBClient
from ....connectors.scrapers import GenericScraper
from ....constants import (
    AuditMetaKeys,
    AuditTypes,
    DBColumns,
    FileTypes,
    FlagResolutions,
    Flags,
    FlagTableColumns,
    MIMETypes,
    SnowFlakeColumns,
)
from ....models import AuditReport
from ....typings import AuditGroupID
from ....utils.collections import group_dicts_by_type
from ....utils.misc import astype_if
from ..report.main import ImageData, PlaceholderDataSRAT, new_report

logger = logger.new_logger(__name__)


class ReportExporter:
    """This class is responsible for exporting audit reports as PDF or PPTX files,
    which includes processing the data and orchestrating the report generation.
    """

    # These flags denote incorrect channel issues.
    flags_incorrect_channel: set[str] = {
        Flags.BAD_TOPIC_CHANNEL,
        Flags.NO_TOPIC_CHANNEL,
    }

    def __init__(self, db: DBClient, config: AuditReport):
        """Initialize the exporter with a database client and a configuration object.

        Args:
            db: Database client.
            config: Configuration object for the report.
        """
        self.db = db
        self.config = config

    @property
    def media_type(self) -> str:
        """Return the media type of the report based on the configuration."""
        return MIMETypes.PDF if self.config.as_pdf else MIMETypes.PPTX

    @property
    def file_extension(self) -> str:
        """Return the file extension of the report based on the configuration."""
        return FileTypes.PDF if self.config.as_pdf else FileTypes.PPTX

    async def export(self) -> tuple[bytes, str, str]:
        """Export the audit report as a PDF or PPTX file.

        Returns:
            Tuple of the exported file content, media type, and file name.
        """
        isrc = SnowFlakeColumns.ISRC
        config = self.config

        actionable_conflict_count: Task[int] = asyncio.create_task(
            self.db.Export.get_actionable_conflict_count(config.audit_group_id)
        )

        actionable_attached_conflict_count: Task[int] = asyncio.create_task(
            self.db.Export.get_actionable_attached_conflict_count(config.audit_group_id)
        )

        task_top_artist_image = (
            asyncio.create_task(self._fetch_top_artist_image_bytes())
            if config.artist_image
            else None
        )

        metadata = await self._fetch_metadata()
        metadata_sr = metadata.get(AuditTypes.SR)

        # Only construct a DF if there are flags, otherwise operations such as
        # getting Series (e.g. flags_sr[isrc]) will raise an error.
        flags = await self._get_qualified_flags(config.audit_group_id)

        if has_flags := flags is not None and not flags.empty:
            flags_sr = flags[flags[DBColumns.TYPE] == AuditTypes.SR]
            flags_at = flags[flags[DBColumns.TYPE] == AuditTypes.AT]

            at_isrcs_flagged_incorrect_channel = flags_at[
                flags_at[FlagTableColumns.FLAG].isin(self.flags_incorrect_channel)
            ]

            at_resolution_starts_with = self._resolution_starts_with(flags_at)

        placeholders = PlaceholderDataSRAT(
            label_name=config.label_name,
            sr_isrcs_unique_total_count=int(
                metadata_sr.get(AuditMetaKeys.FETCHED_ISRCS_UNIQUE_CLEAN)
            ),
            sr_isrcs_unique_corrected=frozenset(
                flags_sr[flags_sr[FlagTableColumns.RESOLUTION].notnull()][isrc].unique()
                if has_flags
                else []
            ),
            sr_isrcs_unique_flagged=frozenset(
                flags_sr[isrc].unique() if has_flags else []
            ),
            sr_isrcs_unique_flagged_yt_ownership_conflict=frozenset(),  # TODO
            sr_isrcs_unique_flagged_ineligible_track=frozenset(),  # TODO
            sr_isrcs_unique_flagged_mrr_ownership_conflict=frozenset(),  # TODO
            sr_isrcs_unique_corrected_yt_ownership_updated=frozenset(),  # TODO
            sr_isrcs_unique_corrected_match_policy_updated=frozenset(),  # TODO
            sr_isrcs_unique_corrected_reference_reactivated=frozenset(),  # TODO
            srugc_avg_daily_views=int(
                metadata_sr.get(AuditMetaKeys.SRUGC_YT_AVG_DAILY_VIEWS)
            ),
            srugc_match_count=int(
                metadata_sr.get(AuditMetaKeys.SRUGC_YT_TOTAL_UGC_CLAIMS)
            ),
            srugc_tracks_with_ugc_match_pct=astype_if(
                metadata_sr.get(AuditMetaKeys.SRUGC_TRACKS_WITH_UGC_MATCH_PCT), float
            ),
            sr_added_territory_rights_count=(
                (
                    at_resolution_starts_with(FlagResolutions.DELIVER_TRACK)
                    | at_resolution_starts_with(FlagResolutions.UPDATE_OWNERSHIP)
                ).sum()
                if has_flags
                else 0
            ),
            sr_actionable_conflict_count=await actionable_conflict_count,
            sr_actionable_attached_conflict_count=await actionable_attached_conflict_count,
            at_isrcs_flagged_yt_ownership_conflict=[],  # TODO
            at_isrcs_corrected_yt_ownership_updated=[],  # TODO
            at_isrcs_total_count=int(
                metadata_sr.get(AuditMetaKeys.FETCHED_ISRCS_CLEAN)
            ),
            at_isrcs_corrected=tuple(
                flags_at[flags_at[FlagTableColumns.RESOLUTION].notnull()][isrc]
                if has_flags
                else []
            ),
            at_isrcs_flagged_yt_incorrect_channel=tuple(
                at_isrcs_flagged_incorrect_channel[isrc] if has_flags else []
            ),
            at_isrcs_corrected_incorrect_channel_fixed=tuple(
                at_isrcs_flagged_incorrect_channel[
                    at_isrcs_flagged_incorrect_channel[
                        FlagTableColumns.RESOLUTION
                    ].notnull()
                ][isrc]
                if has_flags
                else []
            ),
            at_isrcs_flagged=tuple(flags_at[isrc] if has_flags else []),
            at_redelivered_count=(
                sum(at_resolution_starts_with(FlagResolutions.DELIVER_ART_TRACK))
                if has_flags
                else 0
            ),
            at_remapped_count=(
                sum(
                    at_resolution_starts_with(
                        FlagResolutions.LINK_ART_TRACK_TO_CORRECT_TOPIC_CHANNEL
                    )
                )
                if has_flags
                else 0
            ),
        )

        test_images = ImageData(
            top_artist_image=(
                await task_top_artist_image if task_top_artist_image else None
            )
        )

        content = await new_report(
            placeholders,
            test_images,
            template="report_sr_at",
            as_pdf=config.as_pdf,
        )

        file_name = (
            f"{config.label_name} YouTube Audit Analysis"
            f".{self.file_extension.lower()}"
        )

        return content.getvalue(), self.media_type, file_name

    @staticmethod
    def _resolution_starts_with(
        flags_df: pd.DataFrame,
    ) -> Callable[[str], pd.Series]:
        """
        Returns a callable that generates a case-insensitive mask for DataFrame
        rows where the 'resolution' column starts with a specified string,
        accommodating '_REASON' suffixes.

        Args
            flags_df: DataFrame with a 'resolution' column.

        Returns:
            Function that accepts a string and returns a boolean mask for the DataFrame.
        """

        def _(string: str):
            return (
                flags_df[FlagTableColumns.RESOLUTION]
                .str.upper()
                .str.startswith(string.upper())
                .fillna(False)
            )

        return _

    async def _fetch_metadata(self) -> dict[str, dict[str, str]]:
        """Fetch audit group metadata rows and group them by type for easier access."""
        audit_group_meta = await self.db.Meta.get_audit_group(
            self.config.audit_group_id
        )
        return group_dicts_by_type(
            dicts=audit_group_meta,
            type_key=DBColumns.TYPE,
            k_key=DBColumns.KEY,
            v_key=DBColumns.VALUE,
        )

    async def _fetch_top_artist_image_bytes(self) -> BytesIO:
        """Fetch the top artist image as a BytesIO object. The fetched image
        can be either a URL or a base64-encoded image. The function will
        determine the type of the image and fetch it accordingly.
        """
        img_src = self.config.artist_image
        if img_src.startswith("data:image/"):
            # Get the base64 string from the data URL by leaving out the
            # first part of the string (MIME type and encoding).
            _, encoded = img_src.split(",", 1)
            image_data = base64.b64decode(encoded)
        else:
            # Image provided as URL to fetch.
            image_data = await GenericScraper().get(img_src)
        return BytesIO(image_data)

    async def _get_qualified_flags(
        self, audit_group_id: AuditGroupID
    ) -> pd.DataFrame | None:
        """Get the flags that are not ignored for the specified audit group. Those
        include non-resolved flags and resolved flags that are not ignored.

        Reasoning by Ricky Romano (Ops):
            'For any SRs, ATs, MVs where the resolution is Ignore - they should
            not be counted as flagged and they should not be counted as corrected.
            By resolving the flag with Ignore, we are confirming that there is
            nothing wrong with that flagged ISRC for the SR, AT, or MV.'

        Args:
            audit_group_id: Audit group ID.

        Returns:
            DataFrame of qualified flags or None if there are no flags. If API
            returns no flag rows, return None instead of empty DF. We shouldn't
            return an empty DF because it wouldn't have any column names,
            because the returned list of dicts is empty and this would break
            later operations that expect column names to be present.
        """
        if flags := await self.db.AuditGroups.get_flags(audit_group_id):
            flags_df = pd.DataFrame(flags)
            not_ignored_flags = flags_df[
                flags_df[FlagTableColumns.RESOLUTION] != FlagResolutions.IGNORE
            ]
            return None if not_ignored_flags.empty else not_ignored_flags
        return None