src/google_src/gcloud_wrapper.py

"""
Google Cloud wrapper for switching between multiple authenticated accounts.
Supports both Cloud Storage and Google Sheets via gspread.

Credential Sources (in order of priority):
    1. Explicit service account JSON file
    2. ADC (Application Default Credentials) - local or WIF

Usage:
    from google_src.gcloud_wrapper import GCloudWrapper, GCloudAccount
    
    # Initialize with explicit credentials
    gcloud = GCloudWrapper({
        'account1': GCloudAccount('Account 1', credentials_path='./gcloud-key1.json'),
        'account2': GCloudAccount('Account 2', credentials_path='./gcloud-key2.json'),
    })
    
    # Or use ADC (works with local gcloud auth AND Workload Identity Federation)
    gcloud = GCloudWrapper({
        'default': GCloudAccount('Default ADC', use_adc=True),
    })
    
    # Get a storage client
    storage_client = gcloud.get_storage_client('account1')
    bucket = storage_client.bucket('my-bucket')
    
    # Get a gspread client for Sheets
    sheets_client = gcloud.get_sheets_client('account1')
    spreadsheet = sheets_client.open('My Sheet')
    
    # Convenience: get blob directly
    blob = gcloud.get_blob('account1', 'bucket-name', 'path/to/file.txt')
"""

import os
from dataclasses import dataclass, field
from typing import Dict, Optional, List
from contextlib import contextmanager

from google.auth import default as google_auth_default
from google.auth import load_credentials_from_file
from google.auth.credentials import Credentials as BaseCredentials
from google.oauth2 import service_account
from google.cloud import storage
import gspread
from src.config import get_config_value
from src.logger_config import logger



# Default scopes for Google APIs
DEFAULT_SCOPES = [
    "https://www.googleapis.com/auth/cloud-platform",  # Storage
    "https://www.googleapis.com/auth/spreadsheets",    # Sheets
    "https://www.googleapis.com/auth/drive.file",      # Drive (for Sheets)
]


@dataclass
class GCloudAccount:
    """
    Represents a Google Cloud account configuration.
    
    Args:
        name: Human-readable name for the account
        credentials_path: Path to service account JSON file (optional if use_adc=True)
        use_adc: If True, use Application Default Credentials (ADC/WIF)
        project_id: Optional project ID override
        scopes: Optional custom scopes (defaults to storage + sheets + drive)
    """
    name: str
    credentials_path: Optional[str] = None
    use_adc: bool = False
    project_id: Optional[str] = None
    scopes: List[str] = field(default_factory=lambda: DEFAULT_SCOPES.copy())
    
    def __post_init__(self):
        # Validate: must have either credentials_path or use_adc
        if not self.credentials_path and not self.use_adc:
            raise ValueError(
                f"Account '{self.name}': Must provide either credentials_path or set use_adc=True"
            )
        
        # Convert relative path to absolute if needed
        if self.credentials_path:
            if not os.path.isabs(self.credentials_path):
                self.credentials_path = os.path.abspath(self.credentials_path)
            
            # Validate that credentials file exists
            if not os.path.exists(self.credentials_path):
                raise FileNotFoundError(
                    f"Credentials file not found for account '{self.name}': {self.credentials_path}"
                )
    
    def get_credentials(self) -> BaseCredentials:
        """Get Google credentials for this account."""
        if self.use_adc:
            # Use Application Default Credentials (works with local ADC and WIF)
            creds, project = google_auth_default(scopes=self.scopes)
            
            # Capture project from ADC if not explicitly set
            if not self.project_id and project:
                self.project_id = project
            
            # Re-scope if needed (for user credentials from gcloud auth)
            if hasattr(creds, "with_scopes") and not creds.scopes:
                creds = creds.with_scopes(self.scopes)
            
            return creds
        else:
            # Use load_credentials_from_file which supports both:
            # 1. Standard Service Account JSON
            # 2. External Account Credentials (WIF configuration)
            creds, _ = load_credentials_from_file(self.credentials_path, scopes=self.scopes)
            return creds


class GCloudWrapper:
    """
    Wrapper class for managing multiple Google Cloud accounts.
    Provides unified access to Storage and Sheets clients.
    """
    
    def __init__(self, accounts: Optional[Dict[str, GCloudAccount]] = None):
        """
        Initialize the wrapper with optional accounts.
        
        Args:
            accounts: Dictionary mapping account IDs to GCloudAccount objects.
        """
        self._accounts: Dict[str, GCloudAccount] = accounts or {}
        self._credentials_cache: Dict[str, BaseCredentials] = {}
        self._storage_clients: Dict[str, storage.Client] = {}
        self._sheets_clients: Dict[str, gspread.Client] = {}
    
    # ------------------ Account Management ------------------
    
    def add_account(
        self,
        account_id: str,
        name: str,
        credentials_path: Optional[str] = None,
        use_adc: bool = False,
        project_id: Optional[str] = None,
    ) -> None:
        """
        Add a new account to the wrapper.
        
        Args:
            account_id: Unique identifier for this account
            name: Human-readable name for the account
            credentials_path: Path to the JSON credentials file
            use_adc: If True, use ADC instead of explicit credentials
            project_id: Optional project ID override
        """
        self._accounts[account_id] = GCloudAccount(
            name=name,
            credentials_path=credentials_path,
            use_adc=use_adc,
            project_id=project_id,
        )
        # Clear cached clients for this account
        self._clear_account_cache(account_id)
    
    def add_adc_account(
        self,
        account_id: str,
        name: str,
        project_id: Optional[str] = None,
    ) -> None:
        """
        Convenience method to add an ADC-based account.
        
        Args:
            account_id: Unique identifier for this account
            name: Human-readable name for the account
            project_id: Optional project ID override
        """
        self.add_account(
            account_id=account_id,
            name=name,
            use_adc=True,
            project_id=project_id,
        )
    
    def remove_account(self, account_id: str) -> None:
        """Remove an account from the wrapper."""
        if account_id in self._accounts:
            del self._accounts[account_id]
        self._clear_account_cache(account_id)
    
    def list_accounts(self) -> Dict[str, str]:
        """Return a dictionary of account IDs to account names."""
        return {aid: acc.name for aid, acc in self._accounts.items()}
    
    def _clear_account_cache(self, account_id: str) -> None:
        """Clear all cached clients for an account."""
        self._credentials_cache.pop(account_id, None)
        self._storage_clients.pop(account_id, None)
        self._sheets_clients.pop(account_id, None)
    
    def _get_account(self, account_id: str) -> GCloudAccount:
        """Get account by ID with validation."""
        if account_id not in self._accounts:
            raise KeyError(
                f"Account '{account_id}' not found. Available: {list(self._accounts.keys())}"
            )
        return self._accounts[account_id]
    
    def _get_credentials(self, account_id: str) -> BaseCredentials:
        """Get (cached) credentials for an account."""
        if account_id not in self._credentials_cache:
            account = self._get_account(account_id)
            self._credentials_cache[account_id] = account.get_credentials()
        return self._credentials_cache[account_id]
    
    # ------------------ Storage Client ------------------
    
    def get_storage_client(self, account_id: str, force_new: bool = False) -> storage.Client:
        """
        Get a Cloud Storage client for the specified account.
        
        Args:
            account_id: The account ID to get the client for
            force_new: If True, create a new client even if cached
            
        Returns:
            A google.cloud.storage.Client configured for the account
        """
        if force_new or account_id not in self._storage_clients:
            account = self._get_account(account_id)
            creds = self._get_credentials(account_id)
            
            self._storage_clients[account_id] = storage.Client(
                credentials=creds,
                project=account.project_id,
            )
        
        return self._storage_clients[account_id]
    
    # ------------------ Sheets Client ------------------
    
    def get_sheets_client(self, account_id: str, force_new: bool = False) -> gspread.Client:
        """
        Get a gspread client for Google Sheets for the specified account.
        
        Args:
            account_id: The account ID to get the client for
            force_new: If True, create a new client even if cached
            
        Returns:
            A gspread.Client configured for the account
        """
        if force_new or account_id not in self._sheets_clients:
            creds = self._get_credentials(account_id)
            self._sheets_clients[account_id] = gspread.authorize(creds)
        
        return self._sheets_clients[account_id]
    
    # ------------------ Context Manager ------------------
    
    @contextmanager
    def use_account(self, account_id: str):
        """
        Context manager that yields both storage and sheets clients.
        
        Usage:
            with gcloud.use_account('account1') as (storage_client, sheets_client):
                bucket = storage_client.bucket('my-bucket')
                sheet = sheets_client.open('My Sheet')
        """
        storage_client = self.get_storage_client(account_id)
        sheets_client = self.get_sheets_client(account_id)
        yield storage_client, sheets_client
    
    # ------------------ Storage Convenience Methods ------------------
    
    def get_bucket(self, account_id: str, bucket_name: str) -> storage.Bucket:
        """Get a bucket object for the specified account and bucket name."""
        client = self.get_storage_client(account_id)
        return client.bucket(bucket_name)
    
    def get_blob(self, account_id: str, bucket_name: str, blob_path: str) -> storage.Blob:
        """Get a blob object for the specified account, bucket, and path."""
        bucket = self.get_bucket(account_id, bucket_name)
        return bucket.blob(blob_path)
    
    def download_blob(
        self,
        account_id: str,
        bucket_name: str,
        blob_path: str,
        destination_path: str,
    ) -> str:
        """Download a blob to a local file. Returns the destination path."""
        blob = self.get_blob(account_id, bucket_name, blob_path)
        blob.download_to_filename(destination_path)
        return destination_path
    
    def upload_blob(
        self,
        account_id: str,
        bucket_name: str,
        blob_path: str,
        source_path: str,
        content_type: Optional[str] = None,
    ) -> storage.Blob:
        """Upload a local file to a blob. Returns the Blob object."""
        blob = self.get_blob(account_id, bucket_name, blob_path)
        blob.upload_from_filename(source_path, content_type=content_type)
        return blob
    
    def list_blobs(
        self,
        account_id: str,
        bucket_name: str,
        prefix: Optional[str] = None,
    ) -> list:
        """List blobs in a bucket."""
        client = self.get_storage_client(account_id)
        bucket = client.bucket(bucket_name)
        return list(bucket.list_blobs(prefix=prefix))
    
    # ------------------ Sheets Convenience Methods ------------------
    
    def open_spreadsheet(
        self,
        account_id: str,
        sheet_name: Optional[str] = None,
        sheet_id: Optional[str] = None,
    ) -> gspread.Spreadsheet:
        """
        Open a spreadsheet by name or ID.
        
        Args:
            account_id: The account ID
            sheet_name: Name of the spreadsheet (if opening by name)
            sheet_id: ID of the spreadsheet (if opening by ID, preferred)
            
        Returns:
            A gspread.Spreadsheet object
        """
        if not sheet_name and not sheet_id:
            raise ValueError("Must provide either sheet_name or sheet_id")
        
        client = self.get_sheets_client(account_id)
        
        if sheet_id:
            return client.open_by_key(sheet_id)
        return client.open(sheet_name)
    
    def open_worksheet(
        self,
        account_id: str,
        worksheet_name: str,
        sheet_name: Optional[str] = None,
        sheet_id: Optional[str] = None,
    ) -> gspread.Worksheet:
        """
        Open a specific worksheet from a spreadsheet.
        
        Args:
            account_id: The account ID
            worksheet_name: Name of the worksheet tab
            sheet_name: Name of the spreadsheet
            sheet_id: ID of the spreadsheet (preferred)
            
        Returns:
            A gspread.Worksheet object
        """
        spreadsheet = self.open_spreadsheet(account_id, sheet_name, sheet_id)
        return spreadsheet.worksheet(worksheet_name)


# ------------------ Factory Functions ------------------

def create_adc_wrapper(account_id: str = "default", name: str = "Default ADC") -> GCloudWrapper:
    """
    Create a wrapper with a single ADC-based account.
    Works with local gcloud auth AND Workload Identity Federation.
    
    Usage:
        gcloud = create_adc_wrapper()
        storage = gcloud.get_storage_client('default')
        sheets = gcloud.get_sheets_client('default')
    """
    return GCloudWrapper({
        account_id: GCloudAccount(name=name, use_adc=True)
    })


def create_wrapper_from_env(
    env_prefix: str = "GCLOUD_ACCOUNT_",
    default_to_adc: bool = True,
) -> GCloudWrapper:
    """
    Create a GCloudWrapper from environment variables.
    
    Looks for environment variables like:
        GCLOUD_ACCOUNT_1_NAME=Account 1
        GCLOUD_ACCOUNT_1_CREDENTIALS=./gcloud-key1.json
        GCLOUD_ACCOUNT_2_NAME=Account 2
        GCLOUD_ACCOUNT_2_CREDENTIALS=./gcloud-key2.json
    
    Or for ADC-based accounts:
        GCLOUD_ACCOUNT_DEFAULT_NAME=Default
        GCLOUD_ACCOUNT_DEFAULT_ADC=true
    
    Args:
        env_prefix: Prefix for environment variables
        default_to_adc: If True and no accounts found, create a default ADC account
        
    Returns:
        Configured GCloudWrapper instance
    """
    wrapper = GCloudWrapper()
    
    # Find all account identifiers in environment
    account_ids = set()
    for key in os.environ:
        if key.startswith(env_prefix):
            # Extract account ID (e.g., '1' from 'GCLOUD_ACCOUNT_1_NAME')
            parts = key[len(env_prefix):].split("_")
            if parts:
                account_ids.add(parts[0])
    
    # Create accounts
    for aid in sorted(account_ids):
        name_key = f"{env_prefix}{aid}_NAME"
        creds_key = f"{env_prefix}{aid}_CREDENTIALS"
        adc_key = f"{env_prefix}{aid}_ADC"
        project_key = f"{env_prefix}{aid}_PROJECT"
        
        name = os.environ.get(name_key, f"Account {aid}")
        creds = os.environ.get(creds_key)
        use_adc = os.environ.get(adc_key, "").lower() in ("true", "1", "yes")
        project = os.environ.get(project_key)
        
        try:
            if use_adc:
                wrapper.add_adc_account(aid.lower(), name, project_id=project)
            elif creds:
                wrapper.add_account(aid.lower(), name, credentials_path=creds, project_id=project)
        except (FileNotFoundError, ValueError) as e:
            logger.warning(f"Could not add account '{aid}': {e}")
    
    # Default to ADC if no accounts were found
    if not wrapper._accounts and default_to_adc:
        wrapper.add_adc_account("default", "Default ADC")
    
    return wrapper


def create_default_wrapper() -> GCloudWrapper:
    """
    Create default wrapper with two accounts for Elvoro:
    - final_data: Production data (from GCLOUD_FINAL_DATA_CREDENTIALS or ADC)
    - test_data: Test data (from GCLOUD_TEST_DATA_CREDENTIALS or ADC)
    """
    accounts = {}
    
    # Account 1: final_data
    final_creds = get_config_value("GCLOUD_FINAL_DATA_CREDENTIALS")
    final_project = os.environ.get("GCP_PROJECT_FINAL")
    
    if final_creds:
        accounts["final_data"] = GCloudAccount("Final Data", credentials_path=final_creds, project_id=final_project)
    else:
        accounts["final_data"] = GCloudAccount("Final Data (ADC)", use_adc=True, project_id=final_project)
    
    # Account 2: test_data
    test_creds = get_config_value("GCLOUD_TEST_DATA_CREDENTIALS")
    test_project = os.environ.get("GCP_PROJECT_TEST")
    
    if test_creds:
        accounts["test_data"] = GCloudAccount("Test Data", credentials_path=test_creds, project_id=test_project)
    else:
        accounts["test_data"] = GCloudAccount("Test Data (ADC)", use_adc=True, project_id=test_project)
    
    # Alias 'default' to 'final_data' for convenience
    accounts["default"] = accounts["final_data"]
    
    return GCloudWrapper(accounts)


# Singleton instance for reuse across modules
_default_wrapper_instance = None

def get_default_wrapper() -> GCloudWrapper:
    """
    Get the cached default wrapper (singleton).
    Authentication happens only once, subsequent calls return the cached instance.
    """
    global _default_wrapper_instance
    if _default_wrapper_instance is None:
        _default_wrapper_instance = create_default_wrapper()
    return _default_wrapper_instance


if __name__ == "__main__":
    # Example usage
    print("GCloud Wrapper - Example Usage")
    print("-" * 40)
    
    # Create wrapper with ADC (works locally and in CI with WIF)
    gcloud = create_adc_wrapper()
    
    print("Available accounts:", gcloud.list_accounts())
    
    # Get clients
    try:
        storage_client = gcloud.get_storage_client("default")
        print(f"Storage client project: {storage_client.project}")
        
        sheets_client = gcloud.get_sheets_client("default")
        print("Sheets client ready!")
        
    except Exception as e:
        print(f"Error: {e}")
        print("\nTo use ADC locally, run:")
        print("  gcloud auth application-default login")