"""Utils about s3"""

from datetime import datetime, timezone

import boto3


def ls_s3_objects(
    s3_client,
    begin_ts: float,
    end_ts: float,
    bucket_name: str,
    metric_name: str,
):
    """List object names from s3 buckets"""
    objects = []
    tmp_ts = begin_ts

    s3_prefix = ""
    keys = {}
    while tmp_ts < end_ts:
        tmp = datetime.fromtimestamp(tmp_ts, tz=timezone.utc)
        tmp_hour_prefix = (
            f"metrics={metric_name}/year={tmp.year}/month={tmp.month:02}/"
            f"day={tmp.day:02}/hour={tmp.hour:02}"
        )
        if tmp_hour_prefix != s3_prefix:
            s3_prefix = tmp_hour_prefix
            print(
                f"list objects from s3, bucket: {bucket_name}, prefix: {tmp_hour_prefix}"
            )
            objects_from_s3 = s3_client.list_objects_v2(
                Bucket=bucket_name, Prefix=tmp_hour_prefix
            )
            if "Contents" not in objects_from_s3:
                continue
            for obj in objects_from_s3["Contents"]:
                last_index = obj["Key"].rfind("/")
                keys[obj["Key"][:last_index]] = 1

        tmp_prefix = f"{tmp_hour_prefix}/minute={tmp.minute:02}"
        if tmp_prefix in keys:
            objects.append(f"s3://{bucket_name}/{tmp_prefix}/*.parquet")

        tmp_ts = tmp_ts + 60

    return objects


def get_assume_role_session(role_arn, session_name):
    sts_client = boto3.client("sts")
    response = sts_client.assume_role(
        RoleArn=role_arn, RoleSessionName=session_name, DurationSeconds=3600
    )
    credentials = response["Credentials"]

    return credentials


if __name__ == "__main__":
    s3_client = boto3.client("s3")
    results = ls_s3_objects(
        s3_client=s3_client,
        begin_ts=1758799201,
        end_ts=1758799501,
        bucket_name="sli-dev-metrics",
        metric_name="reporter_service_level_probe_cluster_total",
    )
    print(results)