"""
Google Storage (GS)
https://cloud.google.com/storage/docs/samples/storage-generate-signed-url-v4#storage_generate_signed_url_v4-python

Requires a service account key json file (obtained in IAM) and credentials file $GOOGLE_APPLICATION_CREDENTIALS.
    -> To obtain credentials, the organization level credentials issuance restriction must be diabled
    -> To disable the restriction, the Admin user must be granted organization policy change permissions

Use gcloud for admin:
    gcloud storage buckets list
    gcloud storage buckets create gs://kardio
    gsutil ls
    gsutil mb gs://kardio-test
    gsutil rb gs://kardio-test
    gsutil ls gs://kardio-test
    echo "Test content" > /tmp/exp.text
    gsutil cp /tmp/exp.text gs://kardio-test/tmp/exp.text
    gsutil ls -r gs://kardio-test/tmp
    gsutil rm gs://kardio-test/tmp/exp.text
    gsutil rm -r gs://kardio-test/tmp
"""
import os
import datetime
import re
import pickle
from google.cloud import storage
from dl2050utils.core import oget
from dl2050utils.env import config_load
from dl2050utils.fs import json_save

class GS:
    """
        Generates a v4 signed URL on GCS for blob upload/download.
        Timeout are measured in seconds, and set to 24h default
    """
    def __init__(self, service, default_location='europe-west1'):
        cfg = config_load(service)
        # Create credentials file from config yml
        key_dict = oget(cfg,['gcloud','gs_key'])
        assert key_dict['type'] == 'service_account'
        credentials_p = './gs-keyfile.json'
        json_save(credentials_p, key_dict)
        # Set the GOOGLE_APPLICATION_CREDENTIALS env var to use the credentials file
        os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = credentials_p
        # Connect
        self.default_location = default_location
        self.gc = storage.Client()

    # ####################################################################################################
    # Admin
    # ####################################################################################################

    def create_bucket(self, bucket_name):
        """
        Creates a new bucket in Google Cloud Storage.
        Args:
            bucket_name (str): The name of the bucket to create.
        Returns:
            int: 0 if the bucket was created successfully, 1 otherwise.
        """
        try:
            bucket = self.gc.bucket(bucket_name)
            bucket.storage_class = 'STANDARD'
            new_bucket = self.gc.create_bucket(bucket, location=self.default_location)
            print(f'Bucket {new_bucket.name} created in {new_bucket.location} with storage class {new_bucket.storage_class}')
            return 0
        except Exception as exc:
            print(f'create_bucket EXCEPTION: {str(exc)}')
            return 1
        
    def remove_bucket(self, bucket_name, force_delete=False):
        """
        Deletes a bucket from Google Cloud Storage.
        Args:
            bucket_name (str): The name of the bucket to delete.
            force_delete (bool, optional): If True, deletes all objects in the bucket. Defaults to False.
        Returns:
            int: 0 if the bucket was deleted successfully, 1 otherwise.
        """
        try:
            bucket = self.gc.bucket(bucket_name)
            if force_delete:
                # Delete all objects in the bucket
                blobs = bucket.list_blobs()
                for blob in blobs:
                    blob.delete()
                print(f"All objects in bucket '{bucket_name}' have been deleted.")
            bucket.delete()
            print(f"Bucket '{bucket_name}' has been deleted.")
            return 0
        except Exception as exc:
            print(f'remove_bucket EXCEPTION: {str(exc)}')
            return 1
        
    def delete_object(self, bucket_name, blob_name):
        """
        Deletes a single object from the bucket, with exact match.
        Args:
            bucket_name (str): Name of the GCS bucket.
            blob_name (str): Name of the blob (object) to delete.
        Returns:
            int: False if the object was deleted successfully, True otherwise.
        """
        try:
            bucket = self.gc.bucket(bucket_name)
            blob = bucket.blob(blob_name)
            blob.delete()
            print(f"Deleted object '{blob_name}' from bucket '{bucket_name}'.")
            return 0
        except Exception as exc:
            print(f'delete_object EXCEPTION: {str(exc)}')
            return 1
        
    def delete_objects(self, bucket_name, pattern=None, prefix=None, suffix=None):
        """
        Deletes multiple objects from a bucket matching a pattern.
        Args:
            bucket_name (str): Name of the GCS bucket.
            pattern (str, optional): Regular expression pattern to match blob names.
            prefix (str, optional): Prefix string that blob names should start with.
            suffix (str, optional): Suffix string that blob names should end with.
        Returns:
            int: The number of objects deleted.
        """
        try:
            bucket = self.gc.bucket(bucket_name)
            # Decide which blobs to list based on prefix
            blobs = bucket.list_blobs(prefix=prefix) if prefix else bucket.list_blobs()
            deleted_count = 0
            pattern_compiled = re.compile(pattern) if pattern else None
            for blob in blobs:
                blob_name = blob.name
                match = True
                if pattern_compiled:
                    if not pattern_compiled.search(blob_name): match = False
                if prefix:
                    if not blob_name.startswith(prefix):
                        match = False
                if suffix:
                    if not blob_name.endswith(suffix): match = False
                if match:
                    blob.delete()
                    # print(f"Deleted object '{blob_name}' from bucket '{bucket_name}'.")
                    deleted_count += 1
            print(f"Deleted {deleted_count} objects from bucket '{bucket_name}'.")
            return deleted_count
        except Exception as exc:
            print(f'delete_objects EXCEPTION: {str(exc)}')
            return 0
        
    def list(self, bucket_name, subdir=''):
        """
        Lists all files in a specified subdirectory of a bucket.
        Args:
            bucket_name (str): Name of the GCS bucket.
            subdir (str, optional): Subdirectory path within the bucket. Defaults to '' (root).
        Returns:
            list: List of blob names in the specified subdirectory.
        """
        try:
            bucket = self.gc.bucket(bucket_name)
            blobs = bucket.list_blobs(prefix=subdir)
            file_list = [blob.name for blob in blobs if not blob.name.endswith('/')]
            print(f"Files in {bucket_name}/{subdir}: {file_list}")
            return file_list
        except Exception as exc:
            print(f'list_files EXCEPTION: {str(exc)}')
            return []
    
    # ###################################################################################################################
    # Memmory Download, Upload
    # ###################################################################################################################

    def upload_mem(self, bucket_name, blob_name, data, content_type='application/octet-stream', use_pickle=False):
        """
        Uploads data from memory to a specified bucket and blob.
        Args:
            bucket_name (str): Name of the GCS bucket.
            blob_name (str): Name of the blob to upload.
            data (str or bytes): Data to upload. Can be a string or bytes.
            content_type (str, optional): MIME type of the data. Defaults to 'application/octet-stream'.
            use_pickle (bool, optional): If True, serializes the data using pickle before uploading. Defaults to False.
        Returns:
            int: 0 if upload is successful, 1 otherwise.
        Examples:
            gs.upload_mem(bucket_name, blob_name, data="Hello, world!", content_type='text/plain')
            gs.upload_mem(bucket_name, blob_name, data=b'\x89PNG\r\n\x1a...', content_type='image/png')
        """
        try:
            if use_pickle:
                data = pickle.dumps(data)
            elif isinstance(data, str):
                data = data.encode('utf-8')
            elif not isinstance(data, bytes):
                raise TypeError("Data must be bytes or string if use_pickle is False.")
            bucket = self.gc.bucket(bucket_name)
            blob = bucket.blob(blob_name)
            blob.upload_from_string(data, content_type=content_type)
            print(f'Uploaded data to {bucket_name}/{blob_name}')
            return 0
        except Exception as exc:
            print(f'upload_mem EXCEPTION: {str(exc)}')
            return 1
    
    def download_mem(self, bucket_name, blob_name, as_string=False, encoding='utf-8', use_pickle=False):
        """
        Downloads a blob from the bucket into memory.
        Args:
            bucket_name (str): Name of the GCS bucket.
            blob_name (str): Name of the blob to download.
            as_string (bool, optional): If True, decodes the data using the specified encoding.
                                        Ignored if use_pickle is True. Defaults to False.
            encoding (str, optional): The encoding to use when decoding bytes to string. Defaults to 'utf-8'.
            use_pickle (bool, optional): If True, deserializes the data using pickle after downloading.
                                        Defaults to False.
        Returns:
            Any: The data from the blob, possibly decoded or deserialized.
        """
        try:
            bucket = self.gc.bucket(bucket_name)
            blob = bucket.blob(blob_name)
            data = blob.download_as_bytes()
            if use_pickle:
                # Deserialize the data using pickle
                data = pickle.loads(data)
            elif as_string:
                # Decode the data using the specified encoding
                data = data.decode(encoding)
            # If neither use_pickle nor as_string is True, return the raw bytes
            return data
        except Exception as exc:
            print(f'download_mem EXCEPTION: {str(exc)}')
            return None
        
    # ###################################################################################################################
    # File Download, Upload
    # ###################################################################################################################
        
    def upload_file(self, bucket_name, blob_name, p, content_type='application/octet-stream'):
        """
        Uploads a local file to a specified bucket and blob.
        Args:
            bucket_name (str): Name of the GCS bucket.
            blob_name (str): Name of the blob to upload.
            p (str): Local path of the file to upload.
            content_type (str, optional): MIME type of the data. Defaults to 'application/octet-stream'.
        Returns:
            int: 0 if upload is successful, 1 otherwise.
        """
        try:
            bucket = self.gc.bucket(bucket_name)
            blob = bucket.blob(blob_name)
            blob.upload_from_filename(p, content_type=content_type)
            print(f'Uploaded {p} to {bucket_name}/{blob_name}')
            return 0
        except Exception as exc:
            print(f'upload_file EXCEPTION: {str(exc)}')
            return 1

    def download_file(self, bucket_name, blob_name, p):
        """
        Downloads a blob from the bucket to a local file.
        Args:
            bucket_name (str): Name of the GCS bucket.
            blob_name (str): Name of the blob to download.
            p (str): Local path to save the downloaded file.
        Returns:
            int: 0 if upload is successful, 1 otherwise.
        """
        try:
            bucket = self.gc.bucket(bucket_name)
            blob = bucket.blob(blob_name)
            blob.download_to_filename(p)
            print(f'Downloaded {bucket_name}/{blob_name} to {p}')
            return 0
        except Exception as exc:
            print(f'download_file EXCEPTION: {str(exc)}')
            return 1

    # ####################################################################################################
    # Signed urls
    # ####################################################################################################
    
    def upload_url(self, bucket_name, blob_name, timeout=15*60, size=None):
        """
        Generates a signed URL for uploading a blob.
        Args:
            bucket_name (str): Name of the GCS bucket.
            blob_name (str): Name of the blob to upload.
            timeout (int, optional): URL expiration time in seconds. Defaults to 15 minutes.
            size (int, optional): Maximum allowed size of the upload in bytes.
        Returns:
            str or None: Signed URL for uploading or None if an error occurs.
        """
        try:
            bucket = self.gc.bucket(bucket_name)
            blob = bucket.blob(blob_name)
            query_parameters = None if size is None else \
                               {'x-goog-content-length-range':f'0,{size}'}
            url = blob.generate_signed_url(
                version='v4',
                expiration=datetime.timedelta(seconds=timeout),
                method='PUT',
                content_type='application/octet-stream',
                query_parameters=query_parameters
            )
            return url
        except Exception as exc:
            print(f'upload_url EXCEPTION: {str(exc)}')
            return None
        
    def download_url(self, bucket_name, blob_name, timeout=24*3600):
        """
        Generates a signed URL for downloading a blob.
        Args:
            bucket_name (str): Name of the GCS bucket.
            blob_name (str): Name of the blob to download.
            timeout (int, optional): URL expiration time in seconds. Defaults to 24 hours.
        Returns:
            str or None: Signed URL for downloading or None if an error occurs.
        """
        try:
            bucket = self.gc.bucket(bucket_name)
            blob = bucket.blob(blob_name)
            url = blob.generate_signed_url(
                version='v4',
                expiration=datetime.timedelta(seconds=timeout),
                method='GET'
            )
            return url
        except Exception as exc:
            print(f'gcs_download_signed_url EXCEPTION: {str(exc)}')
            return None
        
    def urls(self, bucket_name, blob_name, timeout=24*3600, size=None):
        """
        Generates both upload and download signed URLs for a blob.
        Args:
            bucket_name (str): Name of the GCS bucket.
            blob_name (str): Name of the blob.
            timeout (int, optional): URL expiration time in seconds. Defaults to 24 hours.
            size (int, optional): Maximum allowed size of the upload in bytes.
        Returns:
            tuple: (upload_url, download_url)
        """
        return self.upload_url(bucket_name, blob_name, timeout=timeout, size=size),\
               self.download_url(bucket_name, blob_name, timeout=timeout)
    
    # ####################################################################################################
    # download_folder, upload_folder
    # TODO: not tested
    # ####################################################################################################

    def upload_folder(self, local_folder_path, bucket_name, remote_folder_path='', use_pickle=False, as_string=False, encoding='utf-8'):
        """
        Recursively uploads a local folder to the remote storage, preserving the directory structure.
        Args:
            local_folder_path (str): Path to the local folder to upload.
            bucket_name (str): Name of the remote bucket.
            remote_folder_path (str, optional): Remote folder path within the bucket. Defaults to ''.
            use_pickle (bool, optional): Whether to serialize the files using pickle before uploading. Defaults to False.
            as_string (bool, optional): Whether to read files as strings with the specified encoding. Defaults to False.
            encoding (str, optional): Encoding to use when reading files as strings. Defaults to 'utf-8'.
        """
        try:
            for root, dirs, files in os.walk(local_folder_path):
                for file in files:
                    # Full local file path
                    local_file_path = os.path.join(root, file)
                    # Compute relative path to maintain directory structure
                    relative_path = os.path.relpath(local_file_path, local_folder_path)
                    relative_path = relative_path.replace(os.sep, '/')  # Ensure UNIX-style path separators
                    # Construct remote file path
                    remote_file_path = os.path.join(remote_folder_path, relative_path).replace(os.sep, '/')
                    # Upload the file
                    success = self.upload_file(
                        bucket_name=bucket_name,
                        blob_name=remote_file_path,
                        p=local_file_path,
                        use_pickle=use_pickle,
                        as_string=as_string,
                        encoding=encoding
                    )
                    if not success: print(f"Failed to upload {local_file_path}")
            print(f"Uploaded folder '{local_folder_path}' to '{bucket_name}/{remote_folder_path}'")
        except Exception as exc:
            print(f'upload_folder EXCEPTION: {str(exc)}')

    def download_folder(self, bucket_name, remote_folder_path, local_folder_path, use_pickle=False, as_string=False, encoding='utf-8'):
        """
        Recursively downloads a folder from the remote storage to a local directory, preserving the directory structure.
        Args:
            bucket_name (str): Name of the remote bucket.
            remote_folder_path (str): Remote folder path within the bucket.
            local_folder_path (str): Local path where the folder will be downloaded.
            use_pickle (bool, optional): Whether to deserialize the files using pickle after downloading. Defaults to False.
            as_string (bool, optional): Whether to write files as strings with the specified encoding. Defaults to False.
            encoding (str, optional): Encoding to use when writing files as strings. Defaults to 'utf-8'.
        """
        try:
            bucket = self.gc.bucket(bucket_name)
            blobs = bucket.list_blobs(prefix=remote_folder_path)
            for blob in blobs:
                # Skip if the blob is a directory placeholder
                if blob.name.endswith('/'): continue
                # Compute relative path to maintain directory structure
                relative_path = os.path.relpath(blob.name, remote_folder_path)
                relative_path = relative_path.replace(os.sep, '/')
                # Construct local file path
                local_file_path = os.path.join(local_folder_path, relative_path)
                # Ensure local directories exist
                local_dir = os.path.dirname(local_file_path)
                if not os.path.exists(local_dir): os.makedirs(local_dir)
                # Download the file
                result = self.download_file(
                    bucket_name=bucket_name,
                    blob_name=blob.name,
                    p=local_file_path,
                    use_pickle=use_pickle,
                    as_string=as_string,
                    encoding=encoding
                )
                if result != 0: print(f"Failed to download {blob.name}")
            print(f"Downloaded folder '{bucket_name}/{remote_folder_path}' to '{local_folder_path}'")
        except Exception as exc:
            print(f'download_folder EXCEPTION: {str(exc)}')
