from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import boto3
from mypy_boto3_s3 import S3Client


def _client(region: str, profile: str) -> S3Client:
    return boto3.Session(profile_name=profile, region_name=region).client("s3")


def upload_file(
    region: str,
    profile: str,
    bucket: str,
    local_path: str,
    s3_key: str,
) -> str:
    _client(region, profile).upload_file(local_path, bucket, s3_key)
    return f"s3://{bucket}/{s3_key}"


def upload_dir(
    region: str,
    profile: str,
    bucket: str,
    local_dir: str,
    s3_prefix: str,
    on_progress: Callable[[], None] | None = None,
) -> int:
    root = Path(local_dir)
    files = [file for file in root.rglob("*") if file.is_file()]
    if not files:
        return 0

    client = _client(region, profile)
    prefix = s3_prefix.rstrip("/")

    def upload_one(file_path: Path) -> None:
        key = f"{prefix}/{file_path.relative_to(root)}"
        client.upload_file(str(file_path), bucket, key)
        if on_progress:
            on_progress()

    with ThreadPoolExecutor(max_workers=10) as pool:
        futures = [pool.submit(upload_one, file) for file in files]
        for future in as_completed(futures):
            future.result()

    return len(files)