import numpy as np
import torch
import torch.utils.data as torch_data


import logging
import os
import random
import time
from io import BytesIO

import boto3
import pandas as pd

from boto3.s3.transfer import TransferConfig
from botocore import UNSIGNED
from botocore.config import Config
from PIL import Image
from tqdm import tqdm


logger = logging.getLogger(__name__)


class S3ImgDataset(torch_data.Dataset):
    def __init__(self, s3_img_path_list: list[str]):
        super().__init__()

        # https://github.com/pytorch/pytorch/issues/13246#issuecomment-905703662
        self.s3_img_path_list = pd.array(s3_img_path_list, dtype="string")

        self.s3_transfer_config = TransferConfig(multipart_threshold=5 * 1024**3)
        self.s3_client = boto3.client(
            "s3",
            aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"),
            aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"),
            aws_session_token=os.getenv("AWS_SESSION_TOKEN"),
            config=Config(
                region_name=os.getenv("AWS_REGION", "us-west-2"),
                signature_version=UNSIGNED,
            ),
        )

    def __len__(self):
        return len(self.s3_img_path_list)

    def __getitem__(self, idx):
        # each path is like s3://<bucket>/<prefix>
        try:
            s3_img_path = self.s3_img_path_list[idx]
            bucket, prefix = s3_img_path.replace("s3://", "").split("/", 1)

            bytesio = BytesIO()
            self.s3_client.download_fileobj(
                bucket, prefix, bytesio, Config=self.s3_transfer_config
            )
            bytesio.seek(0)

            img = torch.from_numpy(np.array(Image.open(bytesio)))
        except Exception as e:
            logger.error(f"Error loading image {s3_img_path}: {e}")
            return self.__getitem__(random.randint(0, len(self) - 1))

        return img


def benchmark(
    s3_img_path_list: list[str],
    fake_training_time,
    batch_size,
    num_workers,
    num_threads,
    prefetch_factor,
):
    dataset = S3ImgDataset(s3_img_path_list)
    dataloader = torch_data.DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        num_threads=num_threads,
        prefetch_factor=prefetch_factor,
        persistent_workers=True,
        pin_memory=True,
        drop_last=True,  # Ensure consistent batch sizes
    )

    num_warmup_batches, num_test_batches = 20, 100

    for batch_idx, batch in tqdm(enumerate(dataloader)):
        if batch_idx == num_warmup_batches:
            tic = time.time()

        # fake training
        time.sleep(fake_training_time)


        if batch_idx > num_test_batches + num_warmup_batches:
            break

    per_iter_time = (time.time() - tic) / num_test_batches
    return per_iter_time


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--fake_training_time",
        type=float,
        default=1.0,
        help="fake training time in seconds",
    )
    parser.add_argument(
        "--batch_size", type=int, default=256, help="how many objects to load per batch"
    )
    parser.add_argument(
        "--num_workers", type=int, default=8, help="how many workers to use"
    )
    parser.add_argument(
        "--num_threads", type=int, default=32, help="how many threads to use"
    )
    parser.add_argument(
        "--prefetch_factor", type=int, default=16, help="how many batches to prefetch"
    )

    args = parser.parse_args()

    with open("s3_img_path_list.txt", "r") as f:
        s3_img_path_list = [line.strip() for line in f.readlines() if line.strip()]

    # repeat s3_img_path_list 1000 times
    s3_img_path_list = s3_img_path_list * 1000

    per_iter_time = benchmark(
        s3_img_path_list,
        args.fake_training_time,
        args.batch_size,
        args.num_workers,
        args.num_threads,
        args.prefetch_factor,
    )
    print(f"args: {args}")
    print(f"per_iter_time: {per_iter_time:.3f}s")