# SPDX-License-Identifier: PROPRIETARY
# SPDX-FileCopyrightText: Copyright The Geneva Authors

"""Storage utilities for working with cloud object stores."""

import os
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import pyarrow.fs as pa_fs


def get_azure_storage_account() -> str:
    """Get Azure storage account name from environment.

    Prefers AZURE_STORAGE_ACCOUNT_NAME over AZURE_STORAGE_ACCOUNT.
    """
    account = os.environ.get("AZURE_STORAGE_ACCOUNT_NAME") or os.environ.get(
        "AZURE_STORAGE_ACCOUNT"
    )
    if not account:
        raise ValueError(
            "AZURE_STORAGE_ACCOUNT_NAME must be set for az:// URIs. "
            "This is required for Azure Blob Storage access."
        )
    return account


def filesystem_from_uri(uri: str) -> tuple["pa_fs.FileSystem", str]:
    """Create a PyArrow FileSystem from a URI, with Azure support.

    PyArrow's FileSystem.from_uri() doesn't auto-detect az:// URIs,
    so this function handles Azure URIs by creating an AzureFileSystem explicitly.

    For Azure, the storage account name is read from AZURE_STORAGE_ACCOUNT_NAME
    environment variable. Authentication is handled by DefaultAzureCredential
    (supports workload identity, managed identity, Azure CLI, etc.).

    Args:
        uri: Storage URI (e.g., "s3://bucket/path", "gs://bucket/path", "az://container/path")

    Returns:
        Tuple of (FileSystem, path) where path is the path within the filesystem

    Raises:
        ValueError: If Azure URI is used but AZURE_STORAGE_ACCOUNT_NAME is not set
    """
    import pyarrow.fs as fs

    if uri.startswith("az://"):
        account_name = get_azure_storage_account()
        path = uri.removeprefix("az://")
        azure_fs = fs.AzureFileSystem(account_name=account_name)
        return azure_fs, path

    return fs.FileSystem.from_uri(uri)
