"""
This module contains the logic for handling the processing of client reports,
when the Lambda is run in PRODUCER mode.

In this mode, the Lambda will fetch all available label IDs and their corresponding
periods from Snowflake, excluding those that have already been processed and are
stored in S3. It will then send to the SNS topic a message for each
label ID and period ID pair, which will be processed by the subscriber
(this Lambda, but in CONSUMER mode).
"""

import asyncio
import json
import re
from collections import defaultdict
from dataclasses import dataclass
from functools import partial

from common.src import aws
from common.src import logger
from common.src.connectors.snowflake import run_sql
from common.src.enums import AWSPayload
from common.src.typings import OrchLabelId
from . import sql_queries
from .enums import ColumnAliases
from .. import threadpool
from ..typings import ProducerHandlerQueue, PeriodId
from ... import config

logger = logger.new_logger(__name__)

LabelsAndPeriods = dict[OrchLabelId, list[PeriodId]]

# The earliest period to process for a label, if available. If this label
# has older periods available, they will be skipped. Disable this
# line to perform a full back-fill.
EARLIEST_PERIOD_TO_PROCESS: int = 315


def run():
    """Wrapper to run the producer in async mode."""
    asyncio.run(_main())


@dataclass
class QueueItem:
    """Dataclass to represent a queue item."""

    label_id: OrchLabelId
    period_id: PeriodId

    def __str__(self):
        return f"Label ID: {self.label_id}, Period ID: {self.period_id}"


async def _main():
    """Main entry point to run the producer-consumer workflow.

    This function orchestrates the execution by:
    - Fetching data from Snowflake and S3 in parallel.
    - Producing tasks (label ID and period ID pairs) into an internal async queue.
    - Launching a consumer coroutine that processes each item from the queue.

    The queue enables concurrent and efficient data handling between the producer
    and consumer. The consumer is responsible for processing each item and
    publishing messages to the SNS topic.
    """
    queue = ProducerHandlerQueue(asyncio.Queue())
    logger.debug("Producer-Consumer queue created.")

    consumer_task = asyncio.create_task(_consumer(queue))
    logger.debug("Consumer task started.")

    try:
        await _producer(queue)
        logger.debug("Producer task finished.")
        await consumer_task  # Wait for consumer to exit after sentinel
        logger.debug("Consumer task finished.")
    except Exception as e:
        logger.exception("Unhandled exception in main flow: %s", e)
        raise


async def _producer(queue: ProducerHandlerQueue) -> None:
    """Producer function to produce items to the queue.

    Fetches all available label IDs and their corresponding periods from Snowflake
    and passes them to the queue as items. It also checks S3 for existing reports
    to avoid overwriting/reprocessing them.
    """
    logger.info("Producer is starting to fetch data from Snowflake and S3...")

    # Create S3 client outside the loop to prevent cross-thread issues.
    s3_client = aws.new_s3_client(config.AWS_REGION)

    run_exec = partial(asyncio.get_running_loop().run_in_executor, threadpool.executor)
    labels_periods, s3_reports = await asyncio.gather(
        run_exec(lambda: tuple(run_sql(sql_queries.labels_period_ids()))),
        run_exec(lambda: _get_existing_reports_in_s3(s3_client, config.S3_BUCKET)),
    )

    # If the SQL query returns no rows, we can skip the rest of the function
    if not labels_periods:
        logger.info("Producer has found no label IDs to process, stopping.")
        await queue.put(None)  # Sentinel
        return

    # Extract label IDs and available periods from the SQL query result.
    # Those are the periods that are available to be processed.
    col_label_id, col_available_periods = (
        ColumnAliases.LABEL_ID.lower(),
        ColumnAliases.AVAILABLE_PERIODS.lower(),
    )

    available = {
        row[col_label_id]: json.loads(row[col_available_periods])
        for row in labels_periods
    }

    logger.info(
        f"Producer has fetched data from Snowflake and S3. It has"
        f" found {len(available)} label IDs in Snowflake and"
        f" {len(s3_reports)} label IDs in S3."
    )

    # Remove from available periods those that are already in S3
    # (i.e., those that have been processed in the past, therefore
    # NOT OVERWRITING them).
    if available:
        enqueued_item_count, skipped_item_count = await _producer_loop(
            queue, available, s3_reports, EARLIEST_PERIOD_TO_PROCESS
        )
        logger.info(
            f"Producer has produced {enqueued_item_count} items to the queue"
            f" (each item being a label ID and period ID). A total of "
            f"{skipped_item_count} items were skipped because they are already"
            f" in S3 or because they are older than the earliest period to process."
        )

    else:
        logger.info("Producer has nothing to process, stopping.")

    await queue.put(None)  # Sentinel


def _get_existing_reports_in_s3(
    s3_client, bucket_name: str
) -> dict[OrchLabelId, list[PeriodId]]:
    """Read the contents of the report S3 bucket, using pagination to handle
    large datasets.

    The contents will be parsed to extract the label IDs and period IDs from the
    S3 keys.

    Args:
        s3_client: The S3 client to use for reading the bucket contents.
            This is usually a boto3 client.
        bucket_name (str): The name of the S3 bucket to read from. This is
            where the reports are stored.

    Returns:
        list: A dictionary where the keys are label IDs and the values are lists
            of period IDs. This represents the existing reports in the S3 bucket.

    Example:
        >>> _get_existing_reports_in_s3("my-bucket")
        >>> {1: [100, 101], 2: [200]}
    """
    bucket_items = aws.utils.iter_s3_bucket_contents(
        s3_client=s3_client,
        bucket_name=bucket_name,
    )

    found = defaultdict(list)
    append = found.__getitem__
    pattern_match = re.compile(r"(\d+)/(\d+)/").match  # Match label_id/period_id/
    key = AWSPayload.KEY

    for item in bucket_items:
        match = pattern_match(item[key])
        if match:
            label_id, period_id = map(int, match.groups())
            append(label_id).append(period_id)

    return dict(found)


async def _producer_loop(
    queue: ProducerHandlerQueue,
    available_periods: LabelsAndPeriods,
    skip_periods: LabelsAndPeriods,
    earliest_period_to_process: int,
) -> tuple[int, int]:
    """
    Asynchronously produces label-period pairs into the queue,
    skipping any periods listed in skip_periods for each label.

    Args:
        queue (ProducerHandlerQueue): The queue to put items into.
        available_periods (LabelsAndPeriods): Periods available per label.
        skip_periods (LabelsAndPeriods): Periods to skip per label.
        earliest_period_to_process (int): The earliest period to process. If
            provided, any periods before this will be skipped.

    Returns:
        tuple: A tuple containing the number of items added to the queue
            and the number of skipped items (which were not added).
    """
    added, skipped = 0, 0
    for label_id, periods in available_periods.items():
        skip_for_label = set(skip_periods.get(label_id, []))

        if earliest_period_to_process is not None:
            skip_for_label.update(
                period for period in periods if period < earliest_period_to_process
            )

        skipped += len(skip_for_label)
        to_process = set(periods) - skip_for_label
        for period_id in to_process:
            logger.debug(
                f"Producer is adding label ID {label_id} and period ID {period_id}"
                f" to the queue."
            )
            await queue.put(QueueItem(label_id=label_id, period_id=period_id))
            added += 1
    return added, skipped


async def _consumer(queue: ProducerHandlerQueue) -> None:
    """Consumer function to asynchronously consume items from the internal
    queue and send them to the SNS topic.
    """
    sns_client = aws.new_sns_client(config.AWS_REGION)
    _topic_arn = aws.utils.get_sns_topic_arn(sns_client, config.SNS_TOPIC)
    logger.info(f"Consumer is ready. Using SNS topic: {_topic_arn}")

    _run_exec = partial(asyncio.get_running_loop().run_in_executor, threadpool.executor)

    while True:
        item = await queue.get()
        if item is None:
            logger.debug("Consumer received sentinel. Exiting.")
            break

        message = json.dumps(item.__dict__)
        try:
            await _run_exec(
                lambda: sns_client.publish(TopicArn=_topic_arn, Message=message)
            )
            logger.debug(f"Published to SNS: {item}")
        except Exception as e:
            logger.exception(f"Failed to publish SNS message for {item}: {e}")

    logger.info("Consumer has finished dispatching all messages.")