"""
Cloud Storage Abstraction Layer

This module provides a unified interface for interacting with different cloud storage providers,
including AWS S3, Google Cloud Storage, Azure Blob Storage, and local filesystem.
It aims to abstract away the vendor-specific complexities, offering a consistent API
for common object storage operations.
"""

import os
import gzip
import boto3
from os import environ
from concurrent.futures import ThreadPoolExecutor
from io import IOBase, BytesIO
from typing import List, Optional, Dict
from botocore.exceptions import ClientError
# For type hinting S3 resource
from mypy_boto3_s3.service_resource import S3ServiceResource
from google.cloud.storage import Client
from google.cloud.exceptions import NotFound
from azure.storage.blob import BlobServiceClient, BlobClient, ContainerClient
from azure.core.exceptions import ResourceNotFoundError
from azure.identity import DefaultAzureCredential
# Assuming this module handles credential retrieval
from auth.client_auth import get_credentials
from utils.logger_utils import log_error, log_info, log_warning
import shutil

# Define the maximum number of threads for concurrent operations.
# This prevents overwhelming the system or the cloud provider's API with too many requests.
MAX_CONCURRENT_THREADS = 100


class ObjectStore:
    """
    Abstract base class defining the interface for cloud storage operations.

    This class serves as a contract for all concrete storage provider implementations.
    Any new storage provider must implement all methods defined here to ensure
    compatibility with the abstraction layer.
    """

    def list(self, bucket: str, prefix: str = "") -> List[str]:
        """
        Lists all objects within a specified bucket that match a given prefix.

        This method is synchronous and intended for listing objects under a single prefix.

        Args:
            bucket: The name of the target storage bucket.
            prefix: The prefix to filter objects by. Only objects whose keys start
                    with this prefix will be returned. Defaults to an empty string,
                    which lists all objects in the bucket.

        Returns:
            A list of object keys (strings) found in the bucket matching the prefix.
        """
        raise NotImplementedError(
            "Subclasses must implement the 'list' method.")

    def lists(self, bucket: str, prefixes: List[str]) -> List[str]:
        """
        Concurrently lists objects from a specified bucket across multiple prefixes.

        This method leverages a thread pool to perform multiple 'list' operations in parallel,
        significantly speeding up the process when dealing with many prefixes.

        Args:
            bucket: The name of the target storage bucket.
            prefixes: A list of prefixes to search for objects. Each prefix will trigger
                      a separate concurrent listing operation.

        Returns:
            A combined list of all unique object keys (strings) matching any of the provided prefixes.
        """
        # Determine the number of threads to use, capped by MAX_CONCURRENT_THREADS.
        thread_count = min(len(prefixes), MAX_CONCURRENT_THREADS)
        all_keys: List[str] = []

        with ThreadPoolExecutor(thread_count) as executor:
            # Create parallel list operations for each prefix.
            # `executor.map` applies `self.list` to each (bucket, prefix) pair.
            results = executor.map(
                self.list, [bucket] * len(prefixes), prefixes)
            for result in results:
                all_keys.extend(result)
            # Ensure unique and sorted keys for consistency
            return sorted(list(set(all_keys)))

    def get_object(self, bucket: str, key: str) -> Optional[IOBase]:
        """
        Retrieves the content of a single object as a file-like stream.

        This method should return an `IOBase` object (e.g., `BytesIO` or a file handle)
        that allows reading the object's content. It handles cases where the object
        might not exist by returning `None`.

        Args:
            bucket: The name of the target storage bucket.
            key: The unique identifier (key) of the object to retrieve.

        Returns:
            An `IOBase` stream containing the object's content, or `None` if the object
            is not found.
        """
        raise NotImplementedError(
            "Subclasses must implement the 'get_object' method.")

    def get_objects(self, bucket: str, keys: List[str]) -> List[Optional[IOBase]]:
        """
        Concurrently retrieves the content of multiple objects.

        This method parallelizes the retrieval of multiple objects from a bucket,
        using a thread pool for efficiency.

        Args:
            bucket: The name of the target storage bucket.
            keys: A list of object keys to retrieve.

        Returns:
            A list of `IOBase` streams, where each stream corresponds to an object
            from the `keys` list. If an object is not found, its corresponding
            entry in the list will be `None`.
        """
        thread_count = min(len(keys), MAX_CONCURRENT_THREADS)
        with ThreadPoolExecutor(thread_count) as executor:
            # `executor.map` applies `self.get_object` to each (bucket, key) pair.
            return list(executor.map(self.get_object, [bucket] * len(keys), keys))

    def upload_file(self, bucket: str, key: str, local_filepath: str):
        """
        Uploads a local file to the specified storage bucket with the given key.

        Args:
            bucket: The name of the target storage bucket.
            key: The desired unique identifier (key) for the uploaded object in the bucket.
            local_filepath: The path to the local file that needs to be uploaded.
        """
        raise NotImplementedError(
            "Subclasses must implement the 'upload_file' method.")

    def download_file(self, bucket: str, key: str, local_filepath: str) -> bool:
        """
        Downloads an object from the storage bucket to a local file.

        Args:
            bucket: The name of the source storage bucket.
            key: The unique identifier (key) of the object to download.
            local_filepath: The local path where the downloaded file will be saved.

        Returns:
            True if the download was successful, False otherwise (e.g., object not found).
        """
        raise NotImplementedError(
            "Subclasses must implement the 'download_file' method.")


class Local(ObjectStore):
    """
    Implementation of ObjectStore for local filesystem storage.

    This class allows treating a local directory structure as an object store,
    useful for development, testing, or local data processing. It handles
    basic file operations and transparently manages gzip compression/decompression
    based on the file extension.
    """

    def __init__(self, base_path: str):
        """
        Initializes the Local object store with a root directory path.

        All operations will be performed relative to this base path.

        Args:
            base_path: The absolute or relative path to the root directory
                       that will act as the "bucket" for local storage.
        """
        self.base_path = os.path.abspath(base_path)
        # Ensure the base directory exists
        os.makedirs(self.base_path, exist_ok=True)

    def list(self, bucket: str = "", prefix: str = "") -> List[str]:
        """
        Lists files in the local directory structure matching the given prefix.

        The `bucket` parameter is ignored in this implementation as it operates
        on a single base path. The `prefix` is used to filter files.

        Args:
            bucket: Ignored for Local implementation.
            prefix: The prefix to filter files by.

        Returns:
            A list of relative file paths (keys) within the `base_path` that match
            the given prefix.
        """
        full_prefix_path = os.path.join(self.base_path, prefix)
        dir_path = os.path.dirname(full_prefix_path)

        if not os.path.isdir(dir_path):
            return []  # Directory does not exist, no files to list

        all_files: List[str] = []
        for root, _, files in os.walk(dir_path):
            for file in files:
                relative_path = os.path.relpath(
                    os.path.join(root, file), self.base_path)
                if relative_path.startswith(prefix):
                    all_files.append(relative_path)
        return all_files

    def get_object(self, bucket: str = "", key: str = "") -> Optional[IOBase]:
        """
        Opens a local file and returns its content as a file-like stream.

        It automatically decompresses gzip files if the key ends with '.gz'.

        Args:
            bucket: Ignored for Local implementation.
            key: The relative path to the file within the `base_path`.

        Returns:
            An `IOBase` stream (file handle or BytesIO) of the file's content,
            or `None` if the file is not found.
        """
        if not key:
            return None  # No key provided

        full_path = os.path.join(self.base_path, key)
        try:
            if key.endswith(".gz"):
                return gzip.open(full_path, "rb")
            return open(full_path, "rb")
        except FileNotFoundError:
            return None
        except IOError as e:
            log_error(f"Error opening local file {full_path}: {e}")
            return None

    def upload_file(self, bucket: str = "", key: str = "", local_filepath: str = ""):
        """
        Uploads a local file to the specified path within the local store.

        Automatically compresses the file with gzip if the `key` ends with '.gz'
        and the source file is not already gzipped. It creates necessary
        directories.

        Args:
            bucket: Ignored for Local implementation.
            key: The desired relative path for the uploaded file within `base_path`.
            local_filepath: The path to the local file to upload.
        """
        if not key or not local_filepath:
            raise ValueError(
                "Key and local_filepath must be provided for upload.")

        destination_path = os.path.join(self.base_path, key)
        os.makedirs(os.path.dirname(destination_path), exist_ok=True)

        try:
            with open(local_filepath, "rb") as src_file:
                content = src_file.read()

            if key.endswith(".gz") and not local_filepath.endswith(".gz"):
                # Compress content if the destination key is gzipped but source is not.
                compressed_content = gzip.compress(content)
                with open(destination_path, "wb") as dest_file:
                    dest_file.write(compressed_content)
            else:
                # Direct copy or if source is already gzipped
                with open(destination_path, "wb") as dest_file:
                    dest_file.write(content)
        except FileNotFoundError:
            log_error(f"Error: Local file not found at {local_filepath}")
            raise
        except IOError as e:
            log_error(
                f"Error uploading file to local storage {destination_path}: {e}")
            raise

    def download_file(self, bucket: str = "", key: str = "", local_filepath: str = "") -> bool:
        """
        Downloads a file from the local store to a new local file path.

        Automatically decompresses gzip files if the source `key` ends with '.gz'
        and the destination `local_filepath` does not.

        Args:
            bucket: Ignored for Local implementation.
            key: The relative path to the file within `base_path` to download.
            local_filepath: The path where the downloaded file will be saved.

        Returns:
            True if the download was successful, False if the source file was not found.
        """
        if not key or not local_filepath:
            raise ValueError(
                "Key and local_filepath must be provided for download.")

        source_path = os.path.join(self.base_path, key)
        try:
            # Create parent directories for the destination file if they don't exist
            os.makedirs(os.path.dirname(local_filepath), exist_ok=True)

            if key.endswith(".gz") and not local_filepath.endswith(".gz"):
                # Decompress if source is gzipped but destination is not
                with gzip.open(source_path, "rb") as src_file:
                    with open(local_filepath, "wb") as dest_file:
                        dest_file.write(src_file.read())
            else:
                # Direct copy or if both are gzipped/ungzipped
                with open(source_path, "rb") as src_file:
                    with open(local_filepath, "wb") as dest_file:
                        dest_file.write(src_file.read())
            return True
        except FileNotFoundError:
            log_error(f"Local file not found at {source_path}")
            return False
        except IOError as e:
            log_error(
                f"Error downloading file from local storage {source_path} to {local_filepath}: {e}")
            return False


class S3(ObjectStore):
    """
    AWS S3 storage implementation using the boto3 library.

    This class provides methods to interact with Amazon S3 buckets, including
    listing, retrieving, uploading, and downloading objects. It handles
    AWS-specific authentication and error handling.
    """

    def __init__(self):
        """
        Initializes the S3 client.

        It attempts to retrieve AWS credentials using `get_credentials()`.
        Ensure that `auth.client_auth.get_credentials` is properly configured
        to provide AWS access key ID, secret access key, and session token.
        """
        try:
            creds = get_credentials()
            self.s3_resource: S3ServiceResource = boto3.resource(
                "s3",
                aws_access_key_id=creds.get('access_key_id'),
                aws_secret_access_key=creds.get('secret_access_key'),
                aws_session_token=creds.get('session_token')
            )
        except Exception as e:
            log_error(f"Failed to initialize S3 client: {e}")
            raise

    def list(self, bucket: str, prefix: str = "") -> List[str]:
        """
        Lists objects in an S3 bucket with a given prefix.

        Args:
            bucket: The name of the S3 bucket.
            prefix: The prefix to filter objects by.

        Returns:
            A list of object keys (strings) in the S3 bucket matching the prefix.
        """
        try:
            return [obj.key for obj in self.s3_resource.Bucket(bucket).objects.filter(Prefix=prefix)]
        except ClientError as e:
            log_error(
                f"Error listing S3 objects in bucket '{bucket}' with prefix '{prefix}': {e}")
            raise
        except Exception as e:
            log_error(
                f"An unexpected error occurred while listing S3 objects: {e}")
            raise

    def get_object(self, bucket: str, key: str) -> Optional[IOBase]:
        """
        Retrieves an object from S3 and returns its content as a file-like stream.

        Automatically decompresses gzip content if the object key ends with '.gz'.

        Args:
            bucket: The name of the S3 bucket.
            key: The unique identifier (key) of the object to retrieve.

        Returns:
            An `IOBase` stream of the object's content, or `None` if the object
            is not found (404 or NoSuchKey error).
        """
        try:
            s3_object = self.s3_resource.Bucket(bucket).Object(key)
            response_body = s3_object.get()["Body"]
            # If the key ends with .gz, assume it's gzipped and decompress on the fly.
            # boto3's get() might also return a GzipFile if Content-Encoding is 'gzip'.
            if key.endswith(".gz"):
                return gzip.GzipFile(fileobj=response_body)
            return response_body
        except ClientError as e:
            error_code = e.response.get("Error", {}).get("Code")
            if error_code in ("404", "NoSuchKey"):
                return None  # Object not found
            log_error(
                f"Error retrieving S3 object '{key}' from bucket '{bucket}': {e}")
            raise
        except Exception as e:
            log_error(
                f"An unexpected error occurred while getting S3 object '{key}': {e}")
            raise

    def upload_file(self, bucket: str, key: str, local_filepath: str):
        """
        Uploads a local file to an S3 bucket.

        Handles gzip compression: if the destination `key` ends with '.gz' and
        the `local_filepath` does not, the content is gzipped before upload.
        `ContentEncoding` metadata is set appropriately.

        Args:
            bucket: The name of the target S3 bucket.
            key: The desired S3 object key.
            local_filepath: The path to the local file to upload.
        """
        try:
            extra_args = {}
            # Read content from the local file
            with open(local_filepath, "rb") as f:
                content = f.read()

            if key.endswith(".gz") and not local_filepath.endswith(".gz"):
                # If destination is GZ but source is not, compress the content
                content = gzip.compress(content)
                extra_args["ContentEncoding"] = "gzip"
            elif key.endswith(".gz") and local_filepath.endswith(".gz"):
                # If both are GZ, assume source is already compressed
                extra_args["ContentEncoding"] = "gzip"

            # Upload the BytesIO stream directly
            self.s3_resource.Bucket(bucket).upload_fileobj(
                BytesIO(content),
                key,
                ExtraArgs=extra_args
            )
        except FileNotFoundError:
            log_error(f"Error: Local file not found at {local_filepath}")
            raise
        except ClientError as e:
            log_error(
                f"Error uploading S3 object '{key}' to bucket '{bucket}': {e}")
            raise
        except Exception as e:
            log_error(f"An unexpected error occurred during S3 file upload: {e}")
            raise

    def download_file(self, bucket: str, key: str, local_filepath: str) -> bool:
        """
        Downloads an object from S3 to a local file.

        Automatically decompresses gzip files: if the source `key` ends with '.gz'
        and the `local_filepath` does not, the downloaded content is decompressed.

        Args:
            bucket: The name of the source S3 bucket.
            key: The S3 object key to download.
            local_filepath: The local path where the file will be saved.

        Returns:
            True if the download was successful, False if the object was not found.
        """
        # log_info(f"[S3] Starting download of s3://{bucket}/{key} to {local_filepath}")
        
        try:
            # Check if bucket exists and is accessible
            try:
                self.s3_resource.meta.client.head_bucket(Bucket=bucket)
                # log_info(f"[S3] Successfully accessed bucket: {bucket}")
            except ClientError as e:
                error_code = int(e.response['Error']['Code'])
                if error_code == 404:
                    log_error(f"[S3] Bucket not found: {bucket}")
                elif error_code == 403:
                    log_error(f"[S3] Access denied to bucket: {bucket}")
                else:
                    log_error(f"[S3] Error accessing bucket {bucket}: {str(e)}")
                return False
            
            # Check if object exists before downloading
            try:
                self.s3_resource.Object(bucket, key).load()
                # log_info(f"[S3] Object exists: s3://{bucket}/{key}")
            except ClientError as e:
                if e.response['Error']['Code'] == '404':
                    log_error(f"[S3] Object not found: s3://{bucket}/{key}")
                else:
                    log_error(f"[S3] Error checking object existence: {str(e)}")
                return False

            # Ensure the directory for the local file exists
            os.makedirs(os.path.dirname(os.path.abspath(local_filepath)) or '.', exist_ok=True)
            # log_info(f"[S3] Local directory prepared: {os.path.dirname(os.path.abspath(local_filepath)) or '.'}")

            # Download the file directly from S3
            # log_info(f"[S3] Starting download to: {os.path.abspath(local_filepath)}")
            try:
                self.s3_resource.Bucket(bucket).download_file(key, local_filepath)
                # log_info(f"[S3] Successfully downloaded to: {os.path.abspath(local_filepath)}")
                
                # Verify the file was actually written
                if not os.path.exists(local_filepath):
                    log_error(f"[S3] Download completed but file not found at: {os.path.abspath(local_filepath)}")
                    return False
                    
                file_size = os.path.getsize(local_filepath)
                log_info(f"[S3] File size: {file_size} bytes")
                
                if file_size == 0:
                    log_warning(f"[S3] Warning: Downloaded file is empty: {os.path.abspath(local_filepath)}")
                
                # Handle decompression if needed
                if key.endswith(".gz") and not local_filepath.endswith(".gz"):
                    log_info("[S3] Decompressing gzipped file...")
                    try:
                        temp_gzipped_filepath = local_filepath + ".gz"
                        os.rename(local_filepath, temp_gzipped_filepath)
                        
                        with gzip.open(temp_gzipped_filepath, "rb") as src, open(local_filepath, "wb") as dest:
                            shutil.copyfileobj(src, dest)
                            
                        # log_info(f"[S3] Successfully decompressed to: {os.path.abspath(local_filepath)}")
                        os.remove(temp_gzipped_filepath)
                        # log_info("[S3] Removed temporary gzipped file")
                    except Exception as decompress_error:
                        log_error(f"[S3] Error during decompression: {str(decompress_error)}")
                        if os.path.exists(temp_gzipped_filepath):
                            os.rename(temp_gzipped_filepath, local_filepath)
                        return False
                
                return True
                
            except Exception as download_error:
                log_error(f"[S3] Error during download: {str(download_error)}")
                # Clean up potentially corrupted file
                if os.path.exists(local_filepath):
                    try:
                        os.remove(local_filepath)
                        log_info(f"[S3] Removed potentially corrupted file: {local_filepath}")
                    except Exception as e:
                        log_error(f"[S3] Error removing corrupted file: {str(e)}")
                return False
                
        except ClientError as e:
            error_code = e.response.get("Error", {}).get("Code")
            if error_code in ("404", "NoSuchKey"):
                log_error(f"[S3] Object not found: s3://{bucket}/{key}")
                return False
            log_error(f"[S3] Error downloading s3://{bucket}/{key}: {str(e)}")
            log_error(f"[S3] Response: {e.response}")
            return False
            
        except Exception as e:
            log_error(f"[S3] Unexpected error downloading s3://{bucket}/{key}: {str(e)}", exc_info=True)
            return False


class GCS(ObjectStore):
    """
    Google Cloud Storage (GCS) implementation.

    This class interacts with Google Cloud Storage buckets using the `google-cloud-storage`
    library, providing methods for common blob operations. It relies on Application Default
    Credentials for authentication.
    """

    def __init__(self):
        """
        Initializes the Google Cloud Storage client.

        This client automatically handles authentication using Application Default Credentials
        (e.g., GOOGLE_APPLICATION_CREDENTIALS environment variable or gcloud CLI login).
        """
        self.client = Client()

    def list(self, bucket: str, prefix: str = "") -> List[str]:
        """
        Lists blobs in a GCS bucket with a given prefix.

        Args:
            bucket: The name of the GCS bucket.
            prefix: The prefix to filter blobs by.

        Returns:
            A list of blob names (strings) in the GCS bucket matching the prefix.
        """
        try:
            return [blob.name for blob in self.client.list_blobs(bucket, prefix=prefix)]
        except Exception as e:
            log_error(
                f"Error listing GCS blobs in bucket '{bucket}' with prefix '{prefix}': {e}")
            raise

    def get_object(self, bucket: str, key: str) -> Optional[IOBase]:
        """
        Retrieves a blob from GCS and returns its content as a BytesIO stream.

        Args:
            bucket: The name of the GCS bucket.
            key: The unique identifier (key) of the blob to retrieve.

        Returns:
            A `BytesIO` stream of the blob's content, or `None` if the blob
            is not found.
        """
        try:
            blob = self.client.bucket(bucket).blob(key)
            # download_as_bytes() fetches the entire blob into memory.
            # For very large files, consider `download_as_string()` or streaming.
            return BytesIO(blob.download_as_bytes())
        except NotFound:
            return None  # Blob not found
        except Exception as e:
            log_error(
                f"Error retrieving GCS blob '{key}' from bucket '{bucket}': {e}")
            raise

    def upload_file(self, bucket: str, key: str, local_filepath: str):
        """
        Uploads a local file to a GCS bucket.

        Args:
            bucket: The name of the target GCS bucket.
            key: The desired GCS blob name.
            local_filepath: The path to the local file to upload.
        """
        try:
            blob = self.client.bucket(bucket).blob(key)
            blob.upload_from_filename(local_filepath)
        except FileNotFoundError:
            log_error(f"Error: Local file not found at {local_filepath}")
            raise
        except Exception as e:
            log_error(
                f"Error uploading GCS file '{key}' to bucket '{bucket}': {e}")
            raise

    def download_file(self, bucket: str, key: str, local_filepath: str) -> bool:
        """
        Downloads a blob from GCS to a local file.

        Args:
            bucket: The name of the source GCS bucket.
            key: The GCS blob name to download.
            local_filepath: The local path where the file will be saved.

        Returns:
            True if the download was successful, False if the blob was not found.
        """
        try:
            # Ensure the directory for the local file exists
            os.makedirs(os.path.dirname(local_filepath), exist_ok=True)

            blob = self.client.bucket(bucket).blob(key)
            blob.download_to_filename(local_filepath)
            return True
        except NotFound:
            return False  # Blob not found
        except Exception as e:
            log_error(
                f"Error downloading GCS blob '{key}' from bucket '{bucket}': {e}")
            raise


class AzureBlob(ObjectStore):
    """
    Azure Blob Storage implementation.

    This class integrates with Azure Blob Storage using the `azure-storage-blob`
    library. It authenticates using `DefaultAzureCredential`, which supports
    various Azure authentication methods.
    """

    def __init__(self):
        """
        Initializes the Azure Blob Storage client.

        It requires the `STORAGE_ACCOUNT_NAME` environment variable to be set.
        Authentication is handled by `DefaultAzureCredential`, which can use
        environment variables, managed identity, Azure CLI, etc.
        """
        account_name = environ.get("STORAGE_ACCOUNT_NAME")
        if not account_name:
            raise ValueError(
                "Environment variable 'STORAGE_ACCOUNT_NAME' must be set for AzureBlob.")
        try:
            self.client = BlobServiceClient(
                f"https://{account_name}.blob.core.windows.net",
                credential=DefaultAzureCredential()
            )
        except Exception as e:
            log_error(f"Failed to initialize Azure Blob client: {e}")
            raise

    def list(self, bucket: str, prefix: str = "") -> List[str]:
        """
        Lists blobs in an Azure Blob Storage container with a given prefix.

        Args:
            bucket: The name of the Azure Blob container (equivalent to a bucket).
            prefix: The prefix to filter blobs by.

        Returns:
            A list of blob names (strings) in the container matching the prefix.
        """
        try:
            container_client: ContainerClient = self.client.get_container_client(
                bucket)
            return [blob.name for blob in container_client.list_blobs(name_starts_with=prefix)]
        except ResourceNotFoundError:
            log_error(f"Azure container '{bucket}' not found.")
            return []  # Return empty list if container doesn't exist
        except Exception as e:
            log_error(
                f"Error listing Azure blobs in container '{bucket}' with prefix '{prefix}': {e}")
            raise

    def get_object(self, bucket: str, key: str) -> Optional[IOBase]:
        """
        Retrieves an Azure blob and returns its content as a BytesIO stream.

        Args:
            bucket: The name of the Azure Blob container.
            key: The unique identifier (key) of the blob to retrieve.

        Returns:
            A `BytesIO` stream of the blob's content, or `None` if the blob
            is not found.
        """
        try:
            blob_client: BlobClient = self.client.get_blob_client(
                container=bucket, blob=key)
            stream = BytesIO()
            blob_client.download_blob().readinto(stream)
            # Reset stream position to the beginning after reading
            stream.seek(0)
            return stream
        except ResourceNotFoundError:
            log_error(f"Azure blob '{key}' not found in container '{bucket}'.")
            return None  # Blob not found
        except Exception as e:
            log_error(
                f"Error retrieving Azure blob '{key}' from container '{bucket}': {e}")
            raise

    def upload_file(self, bucket: str, key: str, local_filepath: str):
        """
        Uploads a local file to an Azure Blob Storage container.

        Args:
            bucket: The name of the target Azure Blob container.
            key: The desired Azure blob name.
            local_filepath: The path to the local file to upload.
        """
        try:
            blob_client: BlobClient = self.client.get_blob_client(
                container=bucket, blob=key)
            with open(local_filepath, "rb") as data:
                # Overwrite existing blob if it exists
                blob_client.upload_blob(data, overwrite=True)
        except FileNotFoundError:
            log_error(f"Error: Local file not found at {local_filepath}")
            raise
        except Exception as e:
            log_error(
                f"Error uploading Azure blob '{key}' to container '{bucket}': {e}")
            raise

    def download_file(self, bucket: str, key: str, local_filepath: str) -> bool:
        """
        Downloads a blob from Azure Blob Storage to a local file.

        Args:
            bucket: The name of the source Azure Blob container.
            key: The Azure blob name to download.
            local_filepath: The local path where the file will be saved.

        Returns:
            True if the download was successful, False if the blob was not found.
        """
        try:
            # Ensure the directory for the local file exists
            os.makedirs(os.path.dirname(local_filepath), exist_ok=True)

            blob_client: BlobClient = self.client.get_blob_client(
                container=bucket, blob=key)
            with open(local_filepath, "wb") as file:
                blob_client.download_blob().readinto(file)
            return True
        except ResourceNotFoundError:
            log_error(f"Azure blob '{key}' not found in container '{bucket}'.")
            return False  # Blob not found
        except Exception as e:
            log_error(
                f"Error downloading Azure blob '{key}' from container '{bucket}': {e}")
            raise


_object_store_instances: Dict[str, ObjectStore] = {}


def pick_object_store(vendor: str) -> ObjectStore:
    """
    Factory method to create an appropriate ObjectStore client based on the vendor,
    using lazy loading and global caching.

    This function provides a centralized way to instantiate the correct storage
    client implementation (AWS S3, Google Cloud Storage, Azure Blob, or Local)
    without needing to know the specific class name. It ensures that each
    ObjectStore client is initialized only once and then reused for subsequent calls.

    Args:
        vendor: The identifier for the cloud provider. Supported values are
                'aws', 'gcp', 'azure', and 'local' (case-insensitive).

    Returns:
        An instance of a concrete `ObjectStore` implementation (e.g., `S3`, `GCS`, `AzureBlob`, `Local`).

    Raises:
        ValueError: If the provided `vendor` is not supported.
    """
    normalized_vendor = vendor.lower()

    if normalized_vendor not in ["aws", "gcp", "azure", "local"]:
        raise ValueError(
            f"Unsupported cloud provider: '{vendor}'. Supported vendors are 'aws', 'gcp', 'azure', 'local'.")

    if normalized_vendor in _object_store_instances:
        return _object_store_instances[normalized_vendor]

    if normalized_vendor == "local":
        base_path = os.path.join(os.getcwd(), "local_storage")
        instance = Local(base_path=base_path)
    elif normalized_vendor == "aws":
        instance = S3()
    elif normalized_vendor == "gcp":
        instance = GCS()
    elif normalized_vendor == "azure":
        instance = AzureBlob()
    else:
        # This case should ideally not be reached due to the initial check
        raise ValueError(f"Unexpected vendor: {vendor}")

    _object_store_instances[normalized_vendor] = instance
    return instance