import argparse
import boto3
import csv
import os
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 = session.client("s3")
    s3_control = session.client("s3control")

    job_operations = []
    paginator = s3.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:
                        job_operations.append(key)

    if len(job_operations) > 0:
        today = datetime.now().strftime("%d%m%Y%H%M%S")
        manifest_file_path = f"manifest-{today}.csv"
        manifest_s3_key = f"manifests/{manifest_file_path}"

        # Generate the manifest file and upload it to s3 and remove it locally
        with open(manifest_file_path, "w", newline="") as manifest_file:
            writer = csv.writer(manifest_file)
            for key in job_operations:
                writer.writerow([bucket_name, key])

        s3.upload_file(manifest_file_path, bucket_name, manifest_s3_key)
        os.remove(manifest_file_path)

        manifest_etag = s3.head_object(Bucket=bucket_name, Key=manifest_s3_key)[
            "ETag"
        ].strip('"')
        account_id = session.client("sts").get_caller_identity().get("Account")
        response = s3_control.create_job(
            AccountId=account_id,
            ConfirmationRequired=False,
            Operation={
                "S3InitiateRestoreObject": {
                    "ExpirationInDays": keep_days,
                    "GlacierJobTier": tier.upper(),
                }
            },
            Report={
                "Enabled": True,
                "Format": "Report_CSV_20180820",
                "Bucket": f"arn:aws:s3:::{bucket_name}",
                "Prefix": "/s3-batch-restore-job-reports/",
                "ReportScope": "AllTasks",
            },
            Manifest={
                "Spec": {
                    "Format": "S3BatchOperations_CSV_20180820",
                    "Fields": ["Bucket", "Key"],
                },
                "Location": {
                    "ObjectArn": f"arn:aws:s3:::{bucket_name}/manifests/{manifest_file_path}",
                    "ETag": manifest_etag,
                },
            },
            Priority=10,
            RoleArn=f"arn:aws:iam::{account_id}:role/cmn-delphi-batchoperations-role",
        )

        print(f"Job created. Job ID: {response['JobId']}")
    else:
        print("No objects found within the specified date range.")


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()