# ======================================================================================================================
#
# IMPORTS
#
# ======================================================================================================================

import os

import boto3
from loguru import logger

# ======================================================================================================================
#
# CONSTANTS
#
# ======================================================================================================================

S3_URI_SCHEME = "s3://"
S3_BUCKET_SEPARATOR = "/"
S3_BOTO_CLIENT = "s3"
S3_LIST_OBJECTS_V2 = "list_objects_v2"
S3_CONTENTS_KEY = "Contents"
S3_OBJECT_KEY = "Key"

# ======================================================================================================================
#
# FUNCTIONS
#
# ======================================================================================================================


def parse_s3_url(s3_url: str) -> tuple[str, str]:
    """
    :param s3_url: S3 URL to parse and extract the bucket and blob names from.
    :return: Tuple of:
        - Bucket name.
        - Blob name.
    """

    if s3_url.startswith(S3_URI_SCHEME):
        s3_url = s3_url.replace(S3_URI_SCHEME, "")

    parts = s3_url.split(S3_BUCKET_SEPARATOR, 1)
    bucket = parts[0]
    key = parts[1] if len(parts) > 1 else ""

    return bucket, key


def download_s3_file(s3_url: str, local_path: str) -> None:
    """
    :param s3_url: S3 URL to download data from.
    :param local_path: Path to save the contents of the download to.
    """

    bucket, key = parse_s3_url(s3_url)

    s3 = boto3.client(S3_BOTO_CLIENT)

    logger.info(f"Downloading {s3_url} to {local_path}...")
    s3.download_file(bucket, key, local_path)
    logger.success("Done.")


def download_s3_prefix(s3_url: str, local_directory: str) -> None:
    """
    :param s3_url: S3 URL with a prefix to download all objects under (e.g. "s3://bucket/some/prefix/").
    :param local_directory: Local directory to recreate the prefix structure into.
    """

    bucket, prefix = parse_s3_url(s3_url)

    if prefix and not prefix.endswith(S3_BUCKET_SEPARATOR):
        prefix += S3_BUCKET_SEPARATOR

    s3 = boto3.client(S3_BOTO_CLIENT)
    paginator = s3.get_paginator(S3_LIST_OBJECTS_V2)

    logger.info(f"Downloading all objects under s3://{bucket}/{prefix} to {local_directory}...")

    count = 0
    for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
        for obj in page.get(S3_CONTENTS_KEY, []):
            key = obj[S3_OBJECT_KEY]
            relative_path = key[len(prefix) :]

            if not relative_path:
                continue

            local_path = os.path.join(local_directory, relative_path)
            os.makedirs(os.path.dirname(local_path), exist_ok=True)

            s3.download_file(bucket, key, local_path)
            count += 1

    logger.success(f"Downloaded {count} file(s).")
