"""
This module contains the main audit runner. It is responsible for
orchestrating the audit flow.
"""

import asyncio
from functools import partial
from typing import Callable, Sequence

from pandas import DataFrame

from ... import logger
from ...connectors.db import Client
from ...constants import AuditTypes, DBAuditLogActions, DBColumns
from ...logic import audit, look_data, queries
from ...logic.audit import factories
from ...typings import AuditGroupID, AuditID, AuditTypeShort, LabelID, UserID
from ...utils.collections import filter_values

logger = logger.new_logger(__name__)

# This semaphore is used to ensure that only one audit is run at a time, in a
# FIFO manner (current built-in CPython implementation). This prevents an overload
# of audit runs, which could lead to performance issues and lets audit runs be
# nicely executed one after the other.
semaphore_runner = asyncio.Semaphore(1)


async def run(
    user: UserID,
    label_id: LabelID,
    *,
    audit_audio: bool = True,
    audit_art_track: bool = False,
    audit_video: bool = False,
    audit_group_id_callback: asyncio.Future = None,
) -> bool:
    """Run an audit for a given label ID. This is the main entrypoint and orchestrates
    the audit flow.

    Sound recordings are always audited. Art tracks and videos can be optionally
    audited.

    Args:
        label_id (int): Label ID.
        audit_audio (bool, optional): Whether to audit sound recording. Defaults to True.
        audit_art_track (bool, optional): Whether to audit art track. Defaults to False.
        audit_video (bool, optional): Whether to audit video. Defaults to False.
        user (int, optional): User ID. Defaults to 0.
        audit_group_id_callback (asyncio.Future, optional): Callback to pass the
          audit group ID to, if needed. Defaults to None. This allows the caller to
          get the audit group ID after the audit has been started, without having
          to wait for the complete audit run to complete (and instead, will only have
          to wait for the audit group to be created in the DB).

    Returns:
        bool: Whether the audit was successful.
    """

    db, label_name = await asyncio.gather(
        *[factories.db_factory(), queries.get_label_name(label_id)]
    )

    # Create audit group and get its ID
    audit_group_id: AuditGroupID = await db.AuditGroups.new(
        user=user,
        label_id=label_id,
        label_name=label_name,
        include_audio=audit_audio,
        include_video=audit_video,
        include_art_track=audit_art_track,
    )

    if audit_group_id_callback:
        audit_group_id_callback.set_result(audit_group_id)

    # Confirm the audit group ID was created before entering the semaphore
    # zone. This is to provide immediate feedback to the user in the UI and
    # shows him that the audit is running (even though it might not have
    # started yet).
    async with semaphore_runner:
        try:
            # These are the IDs of the individual audits that belong to the audit group
            audit_types = (
                AuditTypes.SR,
                AuditTypes.MV,
                AuditTypes.AT,
            )

            audit_ids: dict[AuditTypes, AuditID] = dict(
                zip(audit_types, await _get_audit_ids(db, audit_group_id, audit_types))
            )

            discard_none = partial(filter_values, values=(None,))

            audit_loggers = discard_none(
                factories.logger_factory(
                    db,
                    user,
                    audit_ids,
                )
            )

            audit_flaggers = discard_none(
                factories.flagger_factory(
                    user,
                    audit_ids,
                )
            )

            audit_meta_setters = discard_none(
                factories.meta_setter_factory(db, audit_ids)
            )

            await asyncio.gather(
                *[
                    audit_logger(DBAuditLogActions.STARTED)
                    for audit_logger in audit_loggers.values()
                ]
            )

            logger.info(
                "Starting YouTube Audit for label ID {}. "
                "Include Sound Recording Audit: {}, "
                "Include Video Audit: {}, "
                "Include Art Track Audit: {}.",
                label_id,
                audit_audio,
                audit_video,
                audit_art_track,
            )

            await asyncio.gather(
                *[
                    audit_logger(DBAuditLogActions.STARTED_FETCH)
                    for audit_logger in audit_loggers.values()
                ]
            )

            base_data: dict[AuditTypeShort, DataFrame] = await look_data.fetch(
                label_id=label_id,
                include_sr=audit_audio,
                include_at=audit_art_track,
                include_mv=audit_video,
            )

            await asyncio.gather(
                *[
                    audit_logger(
                        DBAuditLogActions.COMPLETED_FETCH,
                        row_count=len(base_data[name.lower()]),
                    )
                    for name, audit_logger in audit_loggers.items()
                ]
            )

            # Add a static row index to each DataFrame row. Do not rely on the
            # DataFrame index, as it might be reset during the audit process.
            base_data = {k: _add_static_row_idx(v) for k, v in base_data.items()}

            individual_audits: [AuditTypes, Callable] = {}
            if audit_audio:
                individual_audits[AuditTypes.SR] = audit.audio_audit
            if audit_video:
                individual_audits[AuditTypes.MV] = audit.video_audit
            if audit_art_track:
                individual_audits[AuditTypes.AT] = audit.art_track_audit

            for k, v in individual_audits.items():
                await _individual_analysis(
                    analysis_fn=v,
                    audit_flagger=audit_flaggers[k],
                    audit_logger=audit_loggers[k],
                    audit_meta_setter=audit_meta_setters[k],
                    base_data=base_data,
                    db=db,
                    label_id=label_id,
                    user=user,
                    audit_id=audit_ids[k],
                )

            logger.info("Audit group ID {} run completed successfully.", audit_group_id)

        except Exception as ex:
            ex_str = str(ex)
            logger.exception(
                "Error running audit. Marking as failed. Cause: {}", ex_str
            )
            await db.AuditGroups.mark_run_as_failed(audit_group_id, ex_str)
            return False

    return True


async def _individual_analysis(
    *,
    analysis_fn: Callable,
    audit_flagger: factories.Flagger,
    audit_logger: Callable,
    audit_meta_setter: Callable,
    base_data: dict[AuditTypeShort, DataFrame],
    db: Client,
    label_id: LabelID,
    user: UserID,
    audit_id: AuditID,
):
    """Run an individual analysis on a given audit type (e.g. sound recording,
    music video, art track).

    Args:
        analysis_fn: Analysis function.
        audit_flagger: Audit flagger.
        audit_logger: Audit logger.
        base_data: Base data for the audit.
        db: Database client.
        label_id: Label ID.
        user: ID of the user running the audit.
        audit_id: ID of the audit.
    """
    await audit_logger(DBAuditLogActions.STARTED_ANALYSIS)

    await analysis_fn(
        label_id,
        base_data,
        audit_flagger,
        audit_meta_setter,
        db=db,
        audit_id=audit_id,
    )

    flagger = audit_flagger
    await db.Flags.add(flagger.audit_id, user, flagger.flags)
    logger.info("Added {} flags to audit ID {}.", len(flagger.flags), flagger.audit_id)

    await audit_logger(DBAuditLogActions.COMPLETED_ANALYSIS)
    await audit_logger(DBAuditLogActions.COMPLETED)


def _add_static_row_idx(df: DataFrame) -> DataFrame:
    """Add a static row index to a DataFrame. This is a column with the row index
    that does not change when the DataFrame is manipulated, unlike the default
    DataFrame index, which can be reset during the audit process.

    Args:
        df (DataFrame): DataFrame to add the row index to.

    Returns:
        DataFrame: DataFrame with a static row index.
    """
    if DBColumns.ROW_IDX in df.columns:
        raise ValueError(
            "Column 'row_idx' already exists in DataFrame "
            "but is a reserved column name."
        )
    df[DBColumns.ROW_IDX] = df.index
    return df


async def _get_audit_ids(
    db: Client, audit_group_id: AuditGroupID, audit_types: Sequence[AuditTypes]
) -> tuple[AuditID, ...]:
    """Get audit IDs for a given audit group ID.

    Args:
        db (Client): Database client.
        audit_group_id (int): Audit group ID.
        audit_types (Sequence[AuditTypes]): Audit types to get IDs for.
    """
    get_id = partial(db.Audits.get_id, audit_group_id)
    result = await asyncio.gather(*map(get_id, audit_types))
    return result