#!/usr/bin/env python3
# SPDX-FileCopyrightText: 2025 Ben Cooksley <bcooksley@kde.org>
#
# SPDX-License-Identifier: GPL-2.0-or-later
import os
import re
import sys
import hashlib
import argparse
import minio
import requests

# Helper to split S3 URL components for our images
def split_s3_url_components( image_url ):
    # Break up the image URL into the components we need...
    extraction_re = r"^(?P<host>[a-zA-Z0-9\-\.]+)/(?P<bucket>[a-zA-Z][a-zA-Z0-9\-]+[a-zA-Z0-9])/(?P<prefix>[a-zA-Z0-9\-\/]+)$"
    match = re.match( extraction_re, image_url )
    if not match:
        raise Exception("Unable to determine S3 host, bucket name and prefix from supplied URL (S3 buckets must be in path mode)")

    # Return the now split out components
    return ( match.group('host'), match.group('bucket'), match.group('prefix') )

# Helper to print our status
def log_action( target, action, path ):
    # Skip if we're not in verbose mode
    if not arguments.verbose:
        return

    print(f"[{target}] {action}: {path}")

# Parse the command line arguments we've been given
parser = argparse.ArgumentParser(description='Utility to either sync a given local folder location to/from a given S3 compatible bucket.')
parser.add_argument('--mode', choices=['sync', 'upload', 'download'], help='Determine what sync method should be used')
parser.add_argument('--delete', default=False, action='store_true', help='When in upload or download mode, remove any files from the destination that are not present in the source')
parser.add_argument('--local', type=str, required=True, help='Local path that should be synced to/from')
parser.add_argument('--remote', type=str, required=True, help='URL to remote S3 compatible bucket/prefix (without https://)')
parser.add_argument('--verbose', default=False, action='store_true')
arguments = parser.parse_args()

# Make sure we have not been asked to both upload and download simultaneously as that is not supported
if arguments.mode == 'sync' and arguments.delete:
    print("Unable to delete files when syncing as both sides are destinations")
    sys.exit(1)

# Validate - ensure the local path exists
if not os.path.exists(arguments.local):
    print( f"Local path '{arguments.local}' does not exist")
    sys.exit(1)

# Validate - extract the components of the S3 URL
try:
    s3_host, s3_bucket, s3_prefix = split_s3_url_components( arguments.remote )
except Exception as e:
    print(e)
    sys.exit(1)

# Establish a connection to the S3 bucket now
s3_client = minio.Minio( s3_host, secure=True, cert_check=True )

# If we have an OIDC authentication token then redeem that for credentials
# Unfortunately this requires re-establishing the S3 connection
if 'MINIO_OIDC' in os.environ:
    # Reach out to the KDE Tokens service to redeem our OIDC token for a temporary session token
    token_params = {
        'token': os.getenv('MINIO_OIDC')
    }
    r = requests.post('https://tokens.kde.org/minio/gitlab', data=token_params)
    credentials = r.json()

    # Setup a Minio connection
    print("Connecting using OIDC token")
    s3_client = minio.Minio(
        s3_host,
        secure=True,
        cert_check=True,
        access_key=credentials['AccessKeyId'],
        secret_key=credentials['SecretAccessKey'],
        session_token=credentials['SessionToken']
    )

# Otherwise look to see if an access and secret key are available from the environment
elif 'MINIO_ACCESS_KEY' in os.environ and 'MINIO_SECRET_KEY' in os.environ:
    # Setup a Minio connection
    s3_client = minio.Minio(
        s3_host,
        secure=True,
        cert_check=True,
        access_key=os.environ['MINIO_ACCESS_KEY'],
        secret_key=os.environ['MINIO_SECRET_KEY']
    )

# Make sure the prefix we've been given makes sense to use (ie. ends with a /)
if not s3_prefix.endswith('/'):
    s3_prefix += '/'

# Time to start syncing!
local_files = {}
remote_files = {}

# First build a list of remote files
# This also validates we are able to access the bucket and prefix we have been given
bucket_listing = s3_client.list_objects( s3_bucket, prefix=s3_prefix, recursive=True )
for bucket_file in bucket_listing:
    # Add the file details
    file_path = bucket_file.object_name.removeprefix(s3_prefix)
    remote_files[ file_path ] = bucket_file.etag

# Now we start going over locally present files
# S3 presents the md5sum of the file as it's "etag" which we use to check if a file has changed
for root, directories, files in os.walk( arguments.local ):
    # Determine the appropriate prefix to use for files we're about to publish'
    relative_root = os.path.relpath( root, arguments.local )
    if relative_root == '.':
        relative_root = ''

    # Go over the relevant files
    for filename in files:
        # Determine what the full remote path name will be
        # Because we might be running on Windows, ensure any backlashes are replaced with forward slashes
        object_name = os.path.join( relative_root, filename )
        object_name = object_name.replace('\\', '/')

        # Put together the full local file path
        local_path = os.path.join( root, filename )

        # Calculate the MD5 hash of the file
        hash_generator = hashlib.md5()
        with open(local_path, "rb") as f:
            while file_chunk := f.read(32768):
                hash_generator.update(file_chunk)

        # Register the file in our local list of files
        md5_hash = hash_generator.hexdigest()
        local_files[ object_name ] = md5_hash

# Now that we have the list of files from both sides, it is time to start syncing them up
# Go over the list of remote files first
for entry, remote_file in remote_files.items():
    # Is there anything to do?
    # If the file exists on both sides and has the same hash then go on to the next one
    if entry in local_files and local_files[entry] == remote_file:
        continue

    # Determine the full remote path...
    local_path = os.path.join( arguments.local, entry )
    remote_path = os.path.join( s3_prefix, entry )

    # Do we need to remove the file?
    if arguments.mode == 'upload' and arguments.delete and entry not in local_files:
        # You are the weakest link goodbye
        log_action("Remote", "Removing", entry)
        s3_client.remove_object(s3_bucket, remote_path)
        continue

    # The only thing left to do now is download the remote file, so make sure we are in the right mode
    if arguments.mode not in ['sync', 'download']:
        continue

    # Make sure the folder this is going in exists
    local_parent_path = os.path.join( arguments.local, os.path.dirname(entry) )
    if not os.path.exists( local_parent_path ):
        os.makedirs( local_parent_path )

    # Now we have a folder to put this in, download the file
    log_action("Remote", "Downloading", entry)
    s3_client.fget_object(s3_bucket, remote_path, local_path)

    # Register the newly download file in the local list
    local_files[entry] = remote_file

# Now go over the local files...
for entry, local_file in local_files.items():
    # Is there anything to do?
    # If the file exists on both sides and has the same hash then go on to the next one
    if entry in remote_files and local_file == remote_files[entry]:
        continue

    # Determine the full local path...
    local_path = os.path.join( arguments.local, entry )
    remote_path = os.path.join( s3_prefix, entry )

    # Do we need to remove the file?
    if arguments.mode == 'download' and arguments.delete and entry not in remote_files:
        # You are the weakest link goodbye
        log_action("Local", "Removing", entry)
        os.remove( local_path )
        continue

    # The only thing left to do now is upload the file, so make sure we are in the right mode
    if arguments.mode not in ['sync', 'upload']:
        continue

    # Make it so
    log_action("Local", "Uploading", entry)
    s3_client.fput_object(s3_bucket, remote_path, local_path )