import argparse
import boto3
from datetime import datetime, timedelta
import time
import sys


def parse_args():
    parser = argparse.ArgumentParser(
        description="Generate S3 recovery manifest from Athena inventory",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
NB:
    Awsume backup account credentials for running the script.
    The resulting files can be large, so be careful with disk space.
    The script will download the results to the current working directory.
    If your query does not return any results try using --yesterday as the inventory for today may not be available yet.

Examples:
    # Generate manifest for entire bucket
    python generate_recovery_manifest.py my-bucket-name

    # Generate manifest for specific prefix
    python generate_recovery_manifest.py my-bucket-name --prefix path/to/files/

Output:
    Creates a CSV file named 'recovery_manifest_<bucket_name>.csv' containing:
    - bucket: The S3 bucket name
    - key: The object key/path
    - version_id: The object version ID
        """,
    )
    parser.add_argument(
        "bucket_name", help="Name of the S3 bucket to generate recovery manifest for"
    )
    parser.add_argument(
        "--prefix",
        help='Optional bucket prefix to filter results (e.g., "folder1/subfolder/")',
        default="",
        metavar="PREFIX",
    )
    parser.add_argument(
        "--yesterday",
        action="store_true",
        help="Use yesterday's inventory instead of today's",
    )
    parser.add_argument(
        "--last-modified-newer-than",
        help="Only include objects modified after this date (format: YYYY-MM-DD)",
        metavar="DATE",
    )
    parser.add_argument(
        "--last-modified-older-than",
        help="Only include objects modified before this date (format: YYYY-MM-DD)",
        metavar="DATE",
    )
    return parser.parse_args()


def run_athena_query(
    bucket_name: str,
    prefix: str,
    use_yesterday: bool,
    last_modified_newer_than: str = None,
    last_modified_older_than: str = None,
    size_filter: str = None,
) -> str:
    # Initialize Athena client
    athena = boto3.client("athena", region_name="us-east-2")

    # Set up query parameters
    database = "inventory_db"
    table_name = bucket_name
    output_location = "s3://backup-inventory-temp/collab-script-results/"

    if use_yesterday:
        dt = (datetime.now() - timedelta(days=1)).strftime("%Y-%m-%d-01-00")
    else:
        dt = datetime.now().strftime("%Y-%m-%d-01-00")

    # Construct the query
    where_clause = (
        f"WHERE is_latest = true AND is_delete_marker = false AND dt = '{dt}'"
    )
    if prefix:
        where_clause += f" AND key LIKE '{prefix}%'"
    if last_modified_newer_than:
        where_clause += f" AND last_modified_date >= DATE('{last_modified_newer_than}')"
    if last_modified_older_than:
        where_clause += f" AND last_modified_date < DATE('{last_modified_older_than}')"
    if size_filter:
        where_clause += f" AND {size_filter}"

    query = f"""
    SELECT bucket, key, version_id
    FROM \"{table_name}\"
    {where_clause}
    """

    print(query)

    # Start the query execution
    response = athena.start_query_execution(
        QueryString=query,
        QueryExecutionContext={"Database": database},
        ResultConfiguration={"OutputLocation": output_location},
        ResultReuseConfiguration={
            "ResultReuseByAgeConfiguration": {"Enabled": True, "MaxAgeInMinutes": 60}
        },
    )

    query_execution_id = response["QueryExecutionId"]

    # Wait for query to complete
    while True:
        response = athena.get_query_execution(QueryExecutionId=query_execution_id)
        state = response["QueryExecution"]["Status"]["State"]

        if state in ["SUCCEEDED", "FAILED", "CANCELLED"]:
            break

        time.sleep(1)

    if state != "SUCCEEDED":
        error_message = response["QueryExecution"]["Status"].get(
            "StateChangeReason", "Unknown error"
        )
        raise Exception(f"Query failed: {error_message}")

    return query_execution_id


def download_results(query_execution_id: str, output_file: str):
    # Initialize S3 client
    s3 = boto3.client("s3")

    # Get the query results location
    athena = boto3.client("athena", region_name="us-east-2")
    response = athena.get_query_execution(QueryExecutionId=query_execution_id)
    results_location = response["QueryExecution"]["ResultConfiguration"][
        "OutputLocation"
    ]

    # Parse the S3 URL
    bucket = results_location.split("/")[2]
    key = "/".join(results_location.split("/")[3:])

    # Download the results
    s3.download_file(bucket, key, output_file)

    # Remove the first line from the downloaded file
    with open(output_file, "r") as file:
        lines = file.readlines()

    with open(output_file, "w") as file:
        file.writelines(lines[1:])


def main():
    args = parse_args()

    try:
        print(f"Querying inventory for bucket: {args.bucket_name}")
        if args.prefix:
            print(f"Using prefix filter: {args.prefix}")

        # Generate manifest for files smaller than 5GB
        print("Generating manifest for files smaller than 5GB...")
        small_files_query_id = run_athena_query(
            args.bucket_name,
            args.prefix,
            args.yesterday,
            args.last_modified_newer_than,
            args.last_modified_older_than,
            "size < 5368709120",  # 5GB in bytes
        )

        small_files_output = f"recovery_manifest_{args.bucket_name}_small.csv"
        print(f"Downloading results to: {small_files_output}")
        download_results(small_files_query_id, small_files_output)

        # Generate manifest for files 5GB or larger
        print("Generating manifest for files 5GB or larger...")
        large_files_query_id = run_athena_query(
            args.bucket_name,
            args.prefix,
            args.yesterday,
            args.last_modified_newer_than,
            args.last_modified_older_than,
            "size >= 5368709120",  # 5GB in bytes
        )

        large_files_output = f"recovery_manifest_{args.bucket_name}_large.csv"
        print(f"Downloading results to: {large_files_output}")
        download_results(large_files_query_id, large_files_output)

        print("Recovery manifests generated successfully!")

    except Exception as e:
        print(f"Error: {str(e)}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()