import asyncio
from collections import defaultdict
from typing import Callable

from ..... import logger
from .....connectors.db.db import Client as DBClient
from .....constants import (
    AuditTypes,
    DBColumns,
    FlagResolutions,
    FlagTableColumns,
    RowReportCols,
    RowTypes,
)
from .....typings import AuditGroupID, AuditID, LabelName, RowIndex
from . import sheets, xlsx_builder

logger = logger.new_logger(__name__)


class _Runner:
    _add_keys = {
        RowReportCols.ISSUE,
        RowReportCols.ACTION_REQUIRED,
        RowReportCols.THIRD_PARTY_LABEL,
        RowReportCols.LABEL_COMMENTS,
        RowReportCols.LABEL_COMMENTS_TERRITORIES,
    }

    def __init__(self, db: DBClient, audit_group_id: AuditGroupID):
        self.db = db
        self.audit_group_id = audit_group_id
        self.label_name = None
        self.rows = None
        self.flags = None

    async def run(self) -> bytes:
        """Run the report generation process."""
        audit_group_id = self.audit_group_id
        logger.info("Generating rows report for audit group {}", audit_group_id)
        self.label_name, self.rows, self.flags = await asyncio.gather(
            *[
                self._get_label_name(),
                self.db.Export.rows_table(audit_group_id),
                self.db.AuditGroups.get_flags(audit_group_id),
            ]
        )

        def _io_operations():
            """IO operations to be run in a separate thread."""
            self._handle_rows_resolved_taken_down_as_deleted()
            self._add_report_columns()
            return self._build_report()

        report = await asyncio.to_thread(_io_operations)
        return report

    async def _get_label_name(self) -> LabelName:
        """Get the label name for the audit group."""
        audit_group_data: dict = await self.db.AuditGroups.get(self.audit_group_id)
        return audit_group_data[DBColumns.LABEL_NAME]

    def _audit_group_includes_mv(self) -> bool:
        """Determines if the audit group includes MV data in any of the rows.
        Short-circuits on the first MV row found.

        Returns:
            True if the audit group includes MV data, False otherwise.
        """
        audit_type = DBColumns.TYPE
        type_mv = AuditTypes.MV
        for row in self.rows:
            if row[audit_type] == type_mv:
                return True
        return False

    def _add_report_columns(self) -> None:
        """Add report columns to the rows. Those are columns that might not present
        in the database, but are required for the report. Their content is derived
        from the data fetched from the database.
        """
        add_cols: dict = {
            RowReportCols.LABEL_NAME: self.label_name,
            **{col: None for col in self._add_keys},
        }
        # Make sure any added columns which are already present in the rows are
        # overwritten with the value already present in the rows.
        self.rows = [{**add_cols, **row} for row in self.rows]

    def _handle_rows_resolved_taken_down_as_deleted(self) -> None:
        """Rows with at least one flag with resolution subtype 'TAKEN_DOWN' are
        to be marked as 'NOT_ACTIVE_RELEASE' as row type, so that they
        are not included in the analysis sheet but in the specific 'deleted' sheets.

        This row type change does NOT affect the database, only the report.

        All changes are done in place.
        """
        if not self.flags:
            return

        marked_deleted = self._get_flagged_rows_matching_predicate(
            lambda f: f[FlagTableColumns.RESOLUTION_SUBTYPE]
            == FlagResolutions.TAKEN_DOWN,
        )

        if not marked_deleted:
            return

        for row in self.rows:
            audit_id = row[DBColumns.ID]
            row_idx = row[DBColumns.ROW_IDX]
            if audit_id in marked_deleted and row_idx in marked_deleted[audit_id]:
                row[DBColumns.ROW_TYPE] = RowTypes.NOT_ACTIVE_RELEASE

    def _get_flagged_rows_matching_predicate(
        self, filter_func: Callable[[dict], bool]
    ) -> dict[AuditID, set[RowIndex]]:
        """Check all flags for a given predicate and return a dictionary with audit ID as
        key and set of matching row indexes as value.

        Args:
            filter_func: Predicate function to filter the flags. Accepts a
                dictionary with flag data and returns a boolean.

        Returns:
            Dictionary with audit ID as key and set of matching row indexes as value.
        """
        matching = defaultdict(set)

        for flag in filter(filter_func, self.flags):
            row_idx = flag[FlagTableColumns.ROW_IDX]
            matching[flag[FlagTableColumns.AUDIT]].add(row_idx)

        return dict(matching)

    def _build_report(self) -> bytes:
        """Build the report."""
        rows, flags = self.rows, self.flags

        audit_includes_mv: bool = self._audit_group_includes_mv()

        if audit_includes_mv:
            ordered_sheets_to_add = [
                sheets.AudioNextSteps(rows),
                sheets.VideoNextSteps(rows, flags=flags),
                sheets.VideoMonetizationRestrictions(rows),
                sheets.AudioAudit(rows, flags=flags),
                sheets.ArtTrackAudit(rows, flags=flags),
                sheets.VideoAudit(rows, flags=flags),
                sheets.AudioDeletedCarvedOut(rows),
                sheets.ArtTrackDeletedCarvedOut(rows),
                sheets.VideoPrivateUnlisted(rows, flags=flags),
            ]

        else:
            ordered_sheets_to_add = [
                sheets.AudioNextSteps(rows),
                sheets.AudioAudit(rows, flags=flags),
                sheets.AudioDeletedCarvedOut(rows),
                sheets.ArtTrackAudit(rows, flags=flags),
                sheets.ArtTrackDeletedCarvedOut(rows),
            ]

        report_builder = xlsx_builder.ReportBuilder()
        sheets_with_rows = (sheet for sheet in ordered_sheets_to_add if sheet.rows)

        for sheet in sheets_with_rows:
            report_builder.add_sheet(sheet)

        return report_builder.to_bytes()


async def run(db, audit_group_id: AuditGroupID) -> bytes:
    """Run the rows report generation process. This is the main entry point for the
    rows report generation process.

    Args:
        db: Database client.
        audit_group_id: Audit group ID for which to generate the report.

    Returns:
        XLSX Bytes of the generated report.
    """
    runner = _Runner(db, audit_group_id)
    return await runner.run()