"""S3 utilities with backoff behavior (ported from integration_scripts).

This module ports the full set of S3 helpers used by the SME exporter
with exponential backoff decorators, progress callbacks, key helpers,
and file operations. Behavior mirrors `integration_scripts.s3_backoff_utils`.
"""
from __future__ import annotations

import os
import sys
import math
from typing import Iterable, Generator

from botocore.exceptions import ClientError  # type: ignore

# Ported dependencies
from src.utils.awsretry import AWSRetry
from src.connectors.s3 import (
    s3_client,
    s3_resource,
    S3_CONFIG,
)


class ProgressPercentage:
    """Render simple progress updates during S3 transfer.

    Mirrors the legacy implementation; writes to stdout.
    """

    def __init__(self, o_s3bucket, key_name: str):
        self._key_name = key_name
        boto_client = o_s3bucket.meta.client
        self._size = boto_client.head_object(
            Bucket=o_s3bucket.name, Key=key_name
        )["ContentLength"]
        self._seen_so_far = 0
        sys.stdout.write("\n")
        self.old_check = None

    def __call__(self, bytes_amount: int) -> None:
        self._seen_so_far += bytes_amount
        percentage = (float(self._seen_so_far) / float(self._size)) * 100
        check = math.floor(percentage / 25)

        if self.old_check != check:
            sys.stdout.write(
                "{} {} MB/{} MB ({}%)\n".format(
                    self._key_name,
                    str(round(self._seen_so_far / 1024 / 1024, 2)),
                    str(round(self._size / 1024 / 1024, 2)),
                    str(round(percentage, 2)),
                )
            )
            sys.stdout.flush()
            self.old_check = check


def get_s3_file_key(*args: str) -> str:
    """Join path components with forward slashes for S3 keys."""
    return os.path.join(*args).replace("\\", "/")


@AWSRetry.backoff()
def download_fileobj_backoff(key: str, path: str, bucket_name: str) -> None:
    """Download a key from S3 with backoff and progress callback."""
    bucket_obj = s3_resource.Bucket(bucket_name)
    progress = ProgressPercentage(bucket_obj, key)
    with open(path, "wb") as data:
        bucket_obj.download_fileobj(
            key,
            data,
            Callback=progress,
            Config=S3_CONFIG,
        )


@AWSRetry.backoff()
def copy_from_backoff(
    src_bucket_name: str,
    src_key: str,
    dst_key: str | None = None,
    dst_bucket_name: str | None = None,
    **kwargs,
) -> None:
    """Copy from a key on S3 using S3 Resource with backoff."""
    copy_source = {"Bucket": src_bucket_name, "Key": src_key}
    if not dst_bucket_name and not dst_key:
        return
    if not dst_key:
        dst_key = src_key
    if not dst_bucket_name:
        dst_bucket_name = src_bucket_name
    s3_resource.Object(dst_bucket_name, dst_key).copy_from(
        CopySource=copy_source, **kwargs
    )


@AWSRetry.backoff()
def copy_key_backoff(
    src_bucket_name: str,
    src_key: str,
    dst_bucket_name: str | None = None,
    dst_key: str | None = None,
    **kwargs,
) -> None:
    """Copy a key on S3 using S3 Client with backoff."""
    copy_source = {"Bucket": src_bucket_name, "Key": src_key}
    if not dst_bucket_name and not dst_key:
        return
    if not dst_key:
        dst_key = src_key
    if not dst_bucket_name:
        dst_bucket_name = src_bucket_name
    s3_client.copy(copy_source, dst_bucket_name, dst_key, **kwargs)


@AWSRetry.backoff()
def filter_backoff(key: str, bucket_name: str):
    """Filter S3 keys (prefix) with backoff."""
    bucket_obj = s3_resource.Bucket(bucket_name)
    return bucket_obj.objects.filter(Prefix=key)


@AWSRetry.backoff()
def delete_backoff(key: str, bucket_name: str) -> None:
    """Delete a key with backoff."""
    s3_resource.Object(bucket_name, key).delete()


@AWSRetry.backoff()
def head_object_backoff(key: str, bucket_name: str):
    """HEAD an object with backoff."""
    return s3_client.head_object(Bucket=bucket_name, Key=key)


@AWSRetry.backoff()
def get_object_backoff(key: str, bucket_name: str):
    """GET an object with backoff."""
    return s3_client.get_object(Bucket=bucket_name, Key=key)


@AWSRetry.backoff()
def s3_key_exists(key: str, bucket_name: str) -> bool:
    """Check if an s3 key exists in a given bucket (with backoff)."""
    try:
        head_object_backoff(key, bucket_name)
    except ClientError:
        return False
    else:
        return True


def filter_keys_by_ext(key_list: Iterable[str], file_ext: str) -> list[str]:
    """Return only keys from key_list with extension file_ext."""
    file_ext = file_ext.strip(".")
    file_ext = "." + file_ext
    return [f for f in key_list if os.path.splitext(f)[1] == file_ext]


def filter_file_keys(
    source_key_folder: str, bucket_name: str, file_ext: str | None = None
) -> list[str]:
    """Get all file keys under a folder (prefix)."""
    obj_list = filter_backoff(source_key_folder, bucket_name)
    files_only = [os.path.split(obj.key)[1] for obj in obj_list]
    files_only = [s for s in files_only if s and "/" not in s]
    if file_ext:
        files_only = filter_keys_by_ext(files_only, file_ext)
    return files_only


def move_all_files_in_folder(
    source_key_folder: str,
    output_key_folder: str,
    bucket_name: str,
    file_ext: str | None = None,
) -> None:
    """Move all files in one 'folder' location to another 'folder'."""
    files_only = filter_file_keys(source_key_folder, bucket_name, file_ext)
    for k in files_only:
        source_key = get_s3_file_key(source_key_folder, k)
        target_key = get_s3_file_key(output_key_folder, k)
        try:
            copy_from_backoff(bucket_name, source_key, target_key)
        except Exception as e:  # pragma: no cover
            raise e
        else:
            delete_backoff(source_key, bucket_name)


def delete_all_files_in_folder(
    source_key_folder: str, bucket_name: str, file_ext: str | None = None
) -> None:
    """Delete all files in a 'folder' (prefix)."""
    files_only = filter_file_keys(source_key_folder, bucket_name, file_ext)
    for k in files_only:
        source_key = get_s3_file_key(source_key_folder, k)
        try:
            delete_backoff(source_key, bucket_name)
        except Exception as e:  # pragma: no cover
            raise e


def delete_file_list(key_list: Iterable[str], bucket_name: str) -> None:
    """Delete all keys provided in key_list."""
    for key in key_list:
        delete_backoff(key, bucket_name)


@AWSRetry.backoff()
def put_string_to_s3(key: str, string: str, bucket_name: str) -> None:
    """Put an arbitrary string to an S3 Object."""
    s3_resource.Object(bucket_name, key).put(Body=string)


@AWSRetry.backoff()
def upload_file_backoff(
    src_file_name: str, s3_path: str, bucket_name: str
) -> None:
    """Upload a single file to S3 with backoff."""
    s3_bucket = s3_resource.Bucket(bucket_name)
    s3_bucket.upload_file(src_file_name, s3_path)


@AWSRetry.backoff()
def upload_files(
    file_name_list: Iterable[str],
    local_path: str,
    s3_path: str,
    bucket_name: str,
) -> None:
    """Upload a list of files to an S3 Bucket (with backoff)."""
    s3_bucket = s3_resource.Bucket(bucket_name)
    print(
        "Uploading {} files to S3 bucket `{}`\n".format(
            len(list(file_name_list)), bucket_name
        )
    )
    for finished_file_name in file_name_list:
        local_file_full_path = os.path.join(
            local_path, os.path.basename(finished_file_name)
        )
        s3_file_key = get_s3_file_key(s3_path, finished_file_name)
        print(
            "Uploading `{}` to `{}`...".format(
                local_file_full_path, s3_file_key
            )
        )
        s3_bucket.upload_file(local_file_full_path, s3_file_key)


def get_matching_s3_objects(
    bucket: str, prefix: str = "", suffix: str = "", delimeter: str = ""
) -> Generator[dict, None, None]:
    """Generate objects in an S3 bucket matching prefix/suffix."""
    paginator = s3_client.get_paginator("list_objects_v2")
    kwargs = {"Bucket": bucket}
    if delimeter:
        kwargs["Delimeter"] = delimeter
    prefixes = (prefix,) if isinstance(prefix, str) else prefix
    for key_prefix in prefixes:
        kwargs["Prefix"] = key_prefix
        for page in paginator.paginate(**kwargs):
            try:
                contents = page["Contents"]
            except KeyError:
                break
            for obj in contents:
                key = obj["Key"]
                if key.endswith(suffix):
                    yield obj


def get_matching_s3_keys(
    bucket: str, prefix: str = "", suffix: str = "", delimeter: str = ""
) -> Generator[str, None, None]:
    """Generate object keys in an S3 bucket matching prefix/suffix."""
    for obj in get_matching_s3_objects(bucket, prefix, suffix, delimeter):
        yield obj["Key"]