"""AWS related utilities."""

import json
from typing import Iterator

from common.src.enums import AWSPayload, AWSPrefixes
from common.src.typings import S3URI

from contextlib import contextmanager


def json_or_sns_payload(event: dict) -> dict | None:
    """Normalizes an AWS Lambda trigger event to JSON format.

    This is useful for AWS Lambdas which can be both triggered by SNS or
    directly via a JSON payload. It checks if the event has the SNS
    schema and if so, it will parse the JSON message from the SNS
    payload. If the event is not an SNS event, it will return the
    original event.

    Args:
        event: The event dictionary, which may contain SNS records.

    Returns:
        If the event is inferred to be an SNS event, it will return the
        parsed JSON message or None if there are no records. Otherwise, it returns the
        original event (i.e., a JSON payload, as it is).
    """
    if len(event) == 1 and AWSPayload.RECORDS in event:  # SNS event schema
        records = event[AWSPayload.RECORDS]
        try:
            event = json.loads(records[0][AWSPayload.SNS][AWSPayload.MESSAGE])
        except IndexError:
            # Handle the edge case where there are no records
            # "None" signals it more explicitly than an empty dict
            return None
    return event


def get_s3_bucket_key_from_uri(uri: S3URI) -> tuple[str, str]:
    """Extracts the bucket name and key from a full S3 URI.

    Args:
        uri: The full S3 URI (e.g., s3://bucket-name/folder/file.txt).

    Example:
        >>> get_s3_bucket_key_from_uri("s3://bucket-name/folder/file.txt")
        >>> ("bucket-name", "folder/file.txt")
    """
    if not uri.startswith(AWSPrefixes.S3):
        raise ValueError(f"Invalid S3 path. Must start with '{AWSPrefixes.S3}'")
    bucket = uri.split("/")[2]
    key = "/".join(uri.split("/")[3:])
    return bucket, key


def get_s3_uri_from_bucket_key(bucket_name: str, key: str) -> S3URI:
    """Returns a full S3 URI from a bucket name and key.

    Args:
        bucket_name: The name of the S3 bucket (e.g., "bucket-name").
        key: The key of the S3 object (e.g., "folder/file.txt").

    Example:
        >>> get_s3_uri_from_bucket_key("bucket-name", "folder/file.txt")
        >>> "s3://bucket-name/folder/file.txt"
    """
    if not bucket_name.strip():
        raise ValueError("Bucket name cannot be empty.")
    if not key.strip():
        raise ValueError("Key cannot be empty.")
    return f"{AWSPrefixes.S3}{bucket_name}/{key}"


def iter_s3_bucket_contents(s3_client, bucket_name: str) -> Iterator[dict]:
    """Read the contents of an S3 bucket, using pagination to handle large datasets.

    This is a memory-efficient way to read the contents of an S3 bucket potentially
    containing thousands of items. It uses pagination to fetch the contents in
    manageable chunks, yielding each item as it is retrieved.

    Args:
        s3_client: The S3 client to use for accessing the bucket.
        bucket_name: The name of the S3 bucket to read from.

    Returns:
        An iterator yielding the contents of the bucket.
    """
    paginator = s3_client.get_paginator("list_objects_v2")
    for page in paginator.paginate(
        Bucket=bucket_name, PaginationConfig={"PageSize": 1000}
    ):
        yield from page.get("Contents", [])


def get_sns_topic_arn(sns_client, topic_name: str) -> str:
    """SNS Publish does not accept topic names, only ARNs. This utility
    function retrieves the ARN for a given topic name.

    Args:
        sns_client: The SNS client to use for retrieving the topic ARN.
        topic_name (str): The name of the SNS topic.

    Returns:
        str: The ARN of the SNS topic.

    Raises:
        ValueError: If no ARN is found for the given topic name.
    """
    response = sns_client.list_topics()
    for topic in response[AWSPayload.TOPICS]:
        topic_arn = topic[AWSPayload.TOPIC_ARN]
        if topic_name in topic_arn.split(":")[-1]:
            return topic_arn
    raise ValueError(f"No ARN found for topic name: {topic_name}")


@contextmanager
def s3_object_body_stream(
    s3_client, bucket: str, key: str, *, delete: bool = False
) -> Iterator[bytes]:
    """
    Context manager to stream the body of an S3 object.

    Args:
        s3_client: Boto3 S3 client.
        bucket (str): Name of the S3 bucket.
        key (str): Key of the S3 object.
        delete (bool, optional): If True, deletes the object after reading.
         Defaults to False.

    Yields:
        Iterator[bytes]: A stream-like object for reading the content of the S3 object.
    """
    body = None
    try:
        obj = s3_client.get_object(Bucket=bucket, Key=key)
        body = obj[AWSPayload.BODY]
        yield body
    finally:
        if delete and body is not None:
            s3_client.delete_object(Bucket=bucket, Key=key)