"""S3 utils."""

import csv
from io import StringIO

import boto3

import config

s3_resource = boto3.resource("s3", region_name=config.AWS_DEFAULT_REGION)


def get_relative_path(table_type: str, folder_type: str, run_id: str = config.CURRENT_RUN_ID, postfix: str = "") -> str:
    """Get S3 path without bucket name.

    Args:
        table_type: Table type.
        folder_type: In or out folder.
        run_id: Run ID.
        postfix: Folder name postfix.

    Returns:
        str: S3 path.

    """
    return f"{config.ENVIRONMENT}/PD/{run_id}/{folder_type}/{table_type}/{postfix}"


def get_full_by_relative(relative_path: str) -> str:
    """Get full path by relative path.

    Args:
        relative_path: Relative S3 path.

    Returns:
        str: Full S3 path.

    """
    return f"s3://{config.S3_BUCKET}/{relative_path}"


def get_relative_by_full(full_path: str) -> str:
    """Get relative path by full path.

    Args:
        full_path: Full S3 path.

    Returns:
        str: Relative S3 path.

    """
    return full_path.replace(f"s3://{config.S3_BUCKET}/", "")


def get_full_path(table_type: str, folder_type: str, run_id: str = config.CURRENT_RUN_ID, postfix: str = "") -> str:
    """Get full S3 path.

    Args:
        table_type: Table type.
        folder_type: In or out folder.
        run_id: Run ID.
        postfix: Folder name postfix.

    Returns:
        str: S3 path.

    """
    return get_full_by_relative(get_relative_path(table_type, folder_type, run_id, postfix))


def get_temp_path(postfix: str = config.CURRENT_RUN_ID):
    """Get temp S3 path.

    Args:
        postfix: Folder name postfix.

    Returns:
        str: S3 path.

    """
    return f"s3://{config.S3_BUCKET}/{config.ENVIRONMENT}/PD/{postfix}/temp/"


def delete_folder_files(folder_path: str):
    """Delete all files in folder.

    Args:
        folder_path: Path prefix.

    """
    bucket = s3_resource.Bucket(config.S3_BUCKET)
    bucket.objects.filter(Prefix=folder_path).delete()


def delete_s3_files_except_csv(s3_path: str):
    """Delete all extra files from S3 to prevent load to MySQL errors.

    Args:
        s3_path: S3 folder path.

    """
    bucket = s3_resource.Bucket(config.S3_BUCKET)
    for obj in bucket.objects.filter(Prefix=s3_path):
        if not obj.key.endswith(".csv"):
            obj.delete()


def delete_temp():
    """Delete temp folder."""
    delete_folder_files(get_relative_by_full(get_temp_path()))


def write_csv(s3_path: str, data: list):
    """Write data to S3 as csv.

    Args:
        s3_path: S3 relative path.
        data: Data as list of rows.

    """
    file = StringIO()
    csv.writer(file).writerows(data)
    s3_resource.Object(config.S3_BUCKET, s3_path).put(Body=file.getvalue())


def copy_files(
    from_table_type: str,
    from_folder_type: str,
    to_table_type,
    to_folder_type,
    run_id: str = config.CURRENT_RUN_ID,
    postfix: str = "",
):
    """Copy all files from one folder to another with file prefix.

    Args:
        from_table_type: Source table type.
        from_folder_type: Source folder type.
        to_table_type: Destination table type.
        to_folder_type: Destination folder type.
        run_id: Job run ID.
        postfix: Destination files prefix (path postfix).

    """
    bucket = s3_resource.Bucket(config.S3_BUCKET)
    from_folder = get_relative_path(from_table_type, from_folder_type, run_id)
    to_folder = get_relative_path(to_table_type, to_folder_type, run_id, postfix=postfix)
    for index, obj in enumerate(bucket.objects.filter(Prefix=from_folder)):
        s3_resource.Object(config.S3_BUCKET, f"{to_folder}_{index}").copy_from(
            CopySource=f"{config.S3_BUCKET}/{obj.key}"
        )