import argparse
import boto3
from botocore.exceptions import ClientError
import re
from datetime import datetime, date
from botocore.client import BaseClient


date_format = "%Y-%m-%d"


def restore_objects(
    session: boto3.Session,
    bucket_name: str,
    prefix: str,
    start_date: date,
    end_date: date,
    keep_days: int,
    tier: str,
) -> None:
    s3_client = session.client("s3")
    paginator = s3_client.get_paginator("list_objects_v2")
    page_iterator = paginator.paginate(Bucket=bucket_name, Prefix=prefix)

    for page in page_iterator:
        if "Contents" in page:
            for obj in page["Contents"]:
                key = obj["Key"]
                date_match = re.search(
                    r"\d{4}-\d{2}-\d{2}", key
                )  # Assumes the date is an 8-digit number in the key
                if date_match:
                    last_modified = date_match.group(0)
                    last_modified_date = datetime.strptime(
                        last_modified, date_format
                    ).date()
                    if start_date <= last_modified_date <= end_date:
                        try:
                            s3_client.restore_object(
                                Bucket=bucket_name,
                                Key=key,
                                RestoreRequest={
                                    "Days": keep_days,
                                    "GlacierJobParameters": {
                                        "Tier": tier,
                                    },
                                },
                            )
                            print(f"Initiated restoration for {key}")
                        except ClientError as e:
                            if (
                                e.response["Error"]["Code"]
                                == "RestoreAlreadyInProgress"
                            ):
                                print(f"Restoration in progress: skipping {key}")
                            else:
                                print(f"Error restoring {key}: {e}")


def main():
    parser = argparse.ArgumentParser(
        description="Restore S3 objects from Glacier Deep Archive within a date range."
    )
    parser.add_argument(
        "--profile", type=str, help="aws cli profile name", default="gdb-delphi-prod"
    )
    parser.add_argument("--bucket", type=str, help="S3 bucket name")
    parser.add_argument("--prefix", type=str, help="S3 prefix")
    parser.add_argument("--start-date", type=str, help="Start date (format: YYYYMMDD)")
    parser.add_argument("--end-date", type=str, help="End date (format: YYYYMMDD)")
    parser.add_argument(
        "--keep-days",
        type=int,
        help="How many days to keep restored objects",
        default=14,
    )
    parser.add_argument(
        "--tier",
        type=str,
        help="Restoration tier, should be Standard, Bulk, or Expedited. Standard by default",
        default="Standard",
    )
    args = parser.parse_args()

    session = boto3.Session(profile_name=args.profile)

    start_date = datetime.strptime(args.start_date, date_format).date()
    end_date = datetime.strptime(args.end_date, date_format).date()

    restore_objects(
        session,
        args.bucket,
        args.prefix,
        start_date,
        end_date,
        args.keep_days,
        args.tier,
    )


if __name__ == "__main__":
    main()