"""
S3 Data Service for Reference Data and Dataset Management

Provides S3-backed storage and caching for reference data and generated datasets.
Handles transparent downloading, uploading, and local caching with validation.

Key Features:
- Reference data download/upload with backup management
- Local caching in ~/.czbillgen/cache/reference_data/
- Basic validation before S3 uploads
- Graceful fallbacks for connectivity issues
- Extensible design for future dataset operations
"""

import json
import os
from datetime import datetime
from pathlib import Path
from typing import Optional, Tuple, List, Dict, Any

import polars as pl
import boto3
from botocore.exceptions import ClientError, NoCredentialsError

from src.log_context import get_logger
from .environment_service import env_service


class S3DataService:
    """
    Service for S3-backed reference data and dataset operations.
    
    Provides centralized S3 operations while maintaining local caching
    for performance and offline capabilities.
    """
    
    def __init__(self):
        """Initialize S3 data service with configuration from environment."""
        self.logger = get_logger(__name__)
        
        # Get S3 configuration from environment service
        s3_config = env_service.get_s3_config()
        self.bucket = s3_config["bucket"]
        self.reference_path = s3_config["reference_path"]
        self.region = s3_config["region"]
        self.profile = s3_config["profile"]
        
        # Local cache configuration
        self.cache_dir = Path.home() / ".czbg" / "cache" / "reference_data"
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        
        # Initialize S3 client
        self._s3_client = None
        
    @property
    def s3_client(self):
        """Lazy initialization of S3 client using boto3 default credential chain."""
        if self._s3_client is None:
            # Use boto3 default credential chain (supports SSO, env vars, IAM roles, etc.)
            # Create session with profile if specified
            if self.profile:
                session = boto3.Session(profile_name=self.profile)
                self._s3_client = session.client('s3', region_name=self.region)
                self.logger.debug(f"S3 client initialized with profile '{self.profile}'")
            else:
                self._s3_client = boto3.client('s3', region_name=self.region)
                self.logger.debug("S3 client initialized using default credential chain")
            
            # Test connection with bucket access
            self._s3_client.head_bucket(Bucket=self.bucket)
                
        return self._s3_client
    
    def _validate_s3_config(self) -> None:
        """Validate that required S3 configuration is present."""
        # Environment service handles all validation and error messages
        env_service.validate_s3_configuration()
    
    def _get_s3_reference_key(self, filename: str) -> str:
        """Get S3 key for reference data file."""
        return f"{self.reference_path}/{filename}"
    
    def _get_cache_file_path(self, filename: str) -> Path:
        """Get local cache path for reference data file."""
        return self.cache_dir / filename
    
    def download_reference_data(self, force: bool = False) -> Path:
        """
        Download reference data from S3 to local cache.
        
        Args:
            force: Force re-download even if cached file exists
            
        Returns:
            Path to local cached reference data file
            
        Raises:
            FileNotFoundError: If reference data not found in S3
            ValueError: If S3 configuration invalid
        """
        self._validate_s3_config()

        cache_file = self._get_cache_file_path("current.parquet")
        s3_key = self._get_s3_reference_key("current.parquet")
        
        # Skip download if file exists and not forcing
        if cache_file.exists() and not force:
            self.logger.info(f"Using cached reference data: {cache_file}")
            return cache_file
        
        self.logger.info(f"Downloading reference data from s3://{self.bucket}/{s3_key}")
        
        # Download from S3
        self.s3_client.download_file(
            Bucket=self.bucket,
            Key=s3_key,
            Filename=str(cache_file)
        )
        
        # Validate downloaded file
        if not self._validate_reference_file(cache_file):
            cache_file.unlink()  # Remove invalid file
            raise ValueError("Downloaded reference data failed validation")
        
        self.logger.info(f"Successfully downloaded reference data to {cache_file}")
        return cache_file
    
    def upload_reference_data(self, local_path: Path, refresh_backup: bool = False) -> None:
        """
        Upload reference data to S3 with optional backup refresh.
        
        Args:
            local_path: Path to local reference data file
            refresh_backup: Whether to refresh backup before uploading
            
        Raises:
            ValueError: If local file fails validation or S3 config invalid
            FileNotFoundError: If local file doesn't exist
        """
        self._validate_s3_config()
        
        if not local_path.exists():
            raise FileNotFoundError(f"Local reference data file not found: {local_path}")
        
        # Validate local file before upload
        if not self._validate_reference_file(local_path):
            raise ValueError(
                f"Local reference data file failed validation: {local_path}. "
                "Check file format, required columns, and row count."
            )
        
        s3_key = self._get_s3_reference_key("reference.parquet")
        backup_key = self._get_s3_reference_key("reference.backup.parquet")
        
        # Refresh backup if requested
        if refresh_backup:
            self._refresh_backup()
        
        # Upload new version
        self.logger.info(f"Uploading reference data to s3://{self.bucket}/{s3_key}")
        
        self.s3_client.upload_file(
            Filename=str(local_path),
            Bucket=self.bucket,
            Key=s3_key
        )
        
        self.logger.info("Successfully uploaded reference data to S3")
    
    def _refresh_backup(self) -> None:
        """Refresh backup by copying current S3 reference data to backup."""
        s3_key = self._get_s3_reference_key("reference.parquet")
        backup_key = self._get_s3_reference_key("reference.backup.parquet")
        
        try:
            # Check if current reference data exists
            self.s3_client.head_object(Bucket=self.bucket, Key=s3_key)
            
            # Copy current to backup
            self.logger.info("Refreshing backup from current reference data")
            
            copy_source = {
                'Bucket': self.bucket,
                'Key': s3_key
            }
            
            self.s3_client.copy_object(
                CopySource=copy_source,
                Bucket=self.bucket,
                Key=backup_key
            )
            
            self.logger.info("Successfully refreshed backup")
            
        except ClientError as e:
            if e.response['Error']['Code'] == '404':
                self.logger.warning("No current reference data found in S3 - skipping backup refresh")
            else:
                self.logger.error(f"Failed to refresh backup: {e}")
                raise
    
    def _validate_reference_file(self, file_path: Path) -> bool:
        """
        Validate reference data file with basic checks.
        
        Args:
            file_path: Path to reference data file
            
        Returns:
            True if file passes validation, False otherwise
        """
        try:
            # Check file exists and not empty
            if not file_path.exists() or file_path.stat().st_size == 0:
                self.logger.error(f"Reference file doesn't exist or is empty: {file_path}")
                return False
            
            # Check valid parquet format
            try:
                df = pl.read_parquet(file_path)
            except Exception as e:
                self.logger.error(f"Invalid parquet format: {e}")
                return False
            
            # Check required columns
            required_columns = {
                "cloud_provider", "service", "usage_family", "resource_type",
                "operation", "usage_type"
            }
            
            missing_columns = required_columns - set(df.columns)
            if missing_columns:
                self.logger.error(f"Missing required columns: {missing_columns}")
                return False
            
            # Check row count sanity
            if df.height < 100:
                self.logger.error(f"Reference data has too few rows: {df.height} (expected >100)")
                return False
            
            self.logger.info(f"Reference file validation passed: {df.height} rows, {len(df.columns)} columns")
            return True
            
        except Exception as e:
            self.logger.error(f"Reference file validation error: {e}")
            return False
    
    def get_cached_reference_path(self) -> Optional[Path]:
        """
        Get path to cached reference data if it exists.
        
        Returns:
            Path to cached file if exists, None otherwise
        """
        cache_file = self._get_cache_file_path("reference.parquet")
        return cache_file if cache_file.exists() else None
    
    def ensure_reference_data(self) -> Path:
        """
        Ensure reference data is available locally, downloading if necessary.
        
        This is the main method for transparent reference data access.
        Downloads from S3 on first access or if cache doesn't exist.
        
        Returns:
            Path to local reference data file
        """
        cached_path = self.get_cached_reference_path()
        
        if cached_path is not None:
            self.logger.debug(f"Using cached reference data: {cached_path}")
            return cached_path
        
        # No cache exists - download from S3
        self.logger.info("No cached reference data found - downloading from S3")
        return self.download_reference_data()
    
    def download_cbf_mapping(self, force: bool = False) -> str:
        """
        Download CBF mapping configuration from S3.
        
        Args:
            force: Force re-download even if cached file exists
            
        Returns:
            CBF mapping content as string
        """
        try:
            # CBF mapping is stored in czbillgen/cbf_mapping/ folder
            s3_key = "czbillgen/cbf_mapping/cbf-mapping.yaml"
            
            # Check cache first (unless force is True)
            cache_file = self.cache_dir / "cbf-mapping.yaml"
            if not force and cache_file.exists():
                self.logger.debug(f"Using cached CBF mapping: {cache_file}")
                return cache_file.read_text()
            
            # Download from S3
            self.logger.info(f"Downloading CBF mapping from s3://{self.bucket}/{s3_key}")
            
            response = self.s3_client.get_object(
                Bucket=self.bucket,
                Key=s3_key
            )
            
            content = response['Body'].read().decode('utf-8')
            
            # Cache the content
            cache_file.write_text(content)
            self.logger.debug(f"CBF mapping cached to: {cache_file}")
            
            return content
            
        except ClientError as e:
            if e.response['Error']['Code'] == 'NoSuchKey':
                self.logger.error(f"CBF mapping not found in S3: {s3_key}")
                raise FileNotFoundError(f"CBF mapping not found in S3: {s3_key}")
            else:
                self.logger.error(f"Failed to download CBF mapping: {e}")
                raise
        except Exception as e:
            self.logger.error(f"Unexpected error downloading CBF mapping: {e}")
            raise

    def download_metadata_profiles(self, force: bool = False) -> str:
        """
        Download metadata profiles configuration from S3.
        
        Args:
            force: Force re-download even if cached file exists
            
        Returns:
            Metadata profiles content as string
        """
        try:
            # Metadata profiles are stored in czbillgen/metadata_profiles/ folder
            s3_key = "czbillgen/metadata_profiles/metadata_profiles.yaml"
            
            # Check cache first (unless force is True)
            cache_file = self.cache_dir / "metadata_profiles.yaml"
            if not force and cache_file.exists():
                self.logger.debug(f"Using cached metadata profiles: {cache_file}")
                return cache_file.read_text()
            
            # Download from S3
            self.logger.info(f"Downloading metadata profiles from s3://{self.bucket}/{s3_key}")
            
            response = self.s3_client.get_object(
                Bucket=self.bucket,
                Key=s3_key
            )
            
            content = response['Body'].read().decode('utf-8')
            
            # Cache the content
            cache_file.write_text(content)
            self.logger.debug(f"Metadata profiles cached to: {cache_file}")
            
            return content
            
        except ClientError as e:
            if e.response['Error']['Code'] == 'NoSuchKey':
                self.logger.error(f"Metadata profiles not found in S3: {s3_key}")
                raise FileNotFoundError(f"Metadata profiles not found in S3: {s3_key}")
            else:
                self.logger.error(f"Failed to download metadata profiles: {e}")
                raise
        except Exception as e:
            self.logger.error(f"Unexpected error downloading metadata profiles: {e}")
            raise

    def file_exists(self, s3_key: str) -> bool:
        """
        Check if a file exists in S3.
        
        Args:
            s3_key: S3 key (path) to check
            
        Returns:
            True if file exists, False otherwise
        """
        try:
            self.s3_client.head_object(Bucket=self.bucket, Key=s3_key)
            return True
        except ClientError as e:
            if e.response['Error']['Code'] == '404':
                return False
            else:
                raise

    def upload_file(self, local_file_path: str, s3_key: str) -> None:
        """
        Upload a file to S3.
        
        Args:
            local_file_path: Path to local file to upload
            s3_key: S3 key (path) to upload to
            
        Raises:
            ValueError: If file doesn't exist or S3 configuration invalid
        """
        self._validate_s3_config()
        
        local_path = Path(local_file_path)
        if not local_path.exists():
            raise ValueError(f"File not found: {local_file_path}")
        
        self.logger.info(f"Uploading file to s3://{self.bucket}/{s3_key}")
        
        self.s3_client.upload_file(
            Filename=str(local_path),
            Bucket=self.bucket,
            Key=s3_key
        )
        
        self.logger.debug(f"Successfully uploaded file: {local_file_path} -> {s3_key}")

    def read_parquet_from_s3(self, s3_key: str) -> pl.DataFrame:
        """
        Read a parquet file from S3 using download workaround.
        
        This method works around Polars S3 credential issues by downloading
        the file to a temporary location and reading it locally.
        
        Args:
            s3_key: S3 key (path) to the parquet file
            
        Returns:
            Polars DataFrame
            
        Raises:
            Exception: If download or read fails
        """
        import tempfile
        import os
        
        with tempfile.NamedTemporaryFile(suffix='.parquet', delete=False) as tmp_file:
            tmp_path = tmp_file.name
        
        try:
            # Download from S3 using boto3
            self.s3_client.download_file(self.bucket, s3_key, tmp_path)
            
            # Read with Polars from local file
            df = pl.read_parquet(tmp_path)
            return df
            
        finally:
            # Clean up temp file
            if os.path.exists(tmp_path):
                os.unlink(tmp_path)

    # Dataset operations for generated billing data
    def _get_dataset_s3_key(self, dataset_name: str, provider: str, month: str) -> str:
        """Get S3 key for dataset file: czbillgen/datasets/{dataset_name}/generate/{provider}/{month}.parquet"""
        return f"czbillgen/datasets/{dataset_name}/generate/{provider}/{month}.parquet"
    
    
    def upload_dataset(self, dataset_name: str, provider: str, month: str, data: pl.DataFrame) -> None:
        """
        Upload dataset to S3 with timestamp/provider/month organization.
        
        Args:
            dataset_name: Human-readable timestamp (e.g., "2025-01-15_14-30-22")
            provider: Cloud provider (e.g., "AWS", "GCP", "Azure")
            month: Month in YYYY-MM format (e.g., "2024-07")
            data: DataFrame to upload
            
        Raises:
            ValueError: If S3 configuration invalid or data is empty
        """
        self._validate_s3_config()
        
        # Basic validation - just ensure data is not empty
        if data.height == 0:
            raise ValueError("Cannot upload empty dataset")
        
        s3_key = self._get_dataset_s3_key(dataset_name, provider, month)
        
        self._upload_dataset_to_key(s3_key, data, dataset_name, provider, month)
    
    def upload_dataset_to_stage(self, dataset_name: str, stage: str, provider: str, month: str, data: pl.DataFrame) -> None:
        """
        Upload dataset to a specific stage in S3.
        
        Args:
            dataset_name: Human-readable timestamp (e.g., "2025-01-15_14-30-22")
            stage: Stage name (e.g., "generate", "metadata", "k8s")
            provider: Cloud provider (e.g., "AWS", "GCP", "Azure")
            month: Month in YYYY-MM format (e.g., "2024-07")
            data: DataFrame to upload
            
        Raises:
            ValueError: If S3 configuration invalid or data is empty
        """
        self._validate_s3_config()
        
        # Basic validation - just ensure data is not empty
        if data.height == 0:
            raise ValueError("Cannot upload empty dataset")
        
        s3_key = f"czbillgen/datasets/{dataset_name}/{stage}/{provider}/{month}.parquet"
        
        self._upload_dataset_to_key(s3_key, data, dataset_name, provider, month)
    
    def _upload_dataset_to_key(self, s3_key: str, data: pl.DataFrame, dataset_name: str, provider: str, month: str) -> None:
        """
        Upload dataset to a specific S3 key.
        
        Args:
            s3_key: S3 key path
            data: DataFrame to upload
            dataset_name: Dataset name for temp file naming
            provider: Provider name for temp file naming
            month: Month for temp file naming
        """
        
        # Write to temporary file first
        temp_file = Path.home() / ".czbg" / "temp" / f"{dataset_name}_{provider}_{month}.parquet"
        temp_file.parent.mkdir(parents=True, exist_ok=True)
        
        try:
            data.write_parquet(temp_file, compression="snappy")
            
            # Upload to S3
            self.logger.info(f"Uploading dataset to s3://{self.bucket}/{s3_key}")
            
            self.s3_client.upload_file(
                Filename=str(temp_file),
                Bucket=self.bucket,
                Key=s3_key
            )
            
            self.logger.info(f"Successfully uploaded dataset: {dataset_name}/{provider}/{month}")
            
        finally:
            # Clean up temp file
            if temp_file.exists():
                temp_file.unlink()
    
    
    def list_datasets(self) -> List[str]:
        """
        List all available datasets in S3.
        
        Returns:
            List of dataset names (timestamps)
        """
        self._validate_s3_config()
        
        response = self.s3_client.list_objects_v2(
            Bucket=self.bucket,
            Prefix="czbillgen/datasets/",
            Delimiter="/"
        )
        
        datasets = []
        if 'CommonPrefixes' in response:
            for prefix in response['CommonPrefixes']:
                # Extract dataset name from prefix: czbillgen/datasets/{dataset_name}/
                dataset_name = prefix['Prefix'].split('/')[-2]
                datasets.append(dataset_name)
        
        return sorted(datasets)
    
    def list_dataset_files(self, dataset_name: str) -> List[Tuple[str, str]]:
        """
        List all files in a specific dataset.
        
        Args:
            dataset_name: Dataset name (timestamp)
            
        Returns:
            List of (provider, month) tuples for available files
        """
        self._validate_s3_config()
        
        prefix = f"czbillgen/datasets/{dataset_name}/generate/"
        response = self.s3_client.list_objects_v2(
            Bucket=self.bucket,
            Prefix=prefix
        )
        
        files = []
        if 'Contents' in response:
            for obj in response['Contents']:
                # Parse key: czbillgen/datasets/{dataset_name}/generate/{provider}/{month}.parquet
                key_parts = obj['Key'].split('/')
                if len(key_parts) >= 6 and key_parts[-1].endswith('.parquet'):
                    provider = key_parts[-2]
                    month = key_parts[-1].replace('.parquet', '')
                    files.append((provider, month))
        
        return sorted(files)
    
    def delete_dataset(self, dataset_name: str) -> None:
        """
        Delete an entire dataset from S3.
        
        Args:
            dataset_name: Dataset name (timestamp)
        """
        self._validate_s3_config()
        
        # List all objects in the dataset
        prefix = f"czbillgen/datasets/{dataset_name}/"
        response = self.s3_client.list_objects_v2(
            Bucket=self.bucket,
            Prefix=prefix
        )
        
        if 'Contents' not in response:
            self.logger.warning(f"Dataset not found: {dataset_name}")
            return
        
        # Delete all objects
        objects_to_delete = [{'Key': obj['Key']} for obj in response['Contents']]
        
        self.s3_client.delete_objects(
            Bucket=self.bucket,
            Delete={'Objects': objects_to_delete}
        )
        
        self.logger.info(f"Successfully deleted dataset: {dataset_name}")

    # Dataset Discovery Methods
    def list_stages(self, dataset_name: str) -> List[str]:
        """
        List all stages for a given dataset.
        
        Args:
            dataset_name: Dataset name to list stages for
            
        Returns:
            List of stage names
        """
        self._validate_s3_config()
        
        prefix = f"czbillgen/datasets/{dataset_name}/"
        response = self.s3_client.list_objects_v2(
            Bucket=self.bucket,
            Prefix=prefix,
            Delimiter="/"
        )
        
        stages = []
        if 'CommonPrefixes' in response:
            for common_prefix in response['CommonPrefixes']:
                # Extract stage from: czbillgen/datasets/{dataset_name}/{stage}/
                stage = common_prefix['Prefix'].split('/')[-2]
                stages.append(stage)
        
        return sorted(stages)

    def list_providers(self, dataset_name: str, stage: str) -> List[str]:
        """
        List all providers for a given dataset+stage.
        
        Args:
            dataset_name: Dataset name
            stage: Stage name
            
        Returns:
            List of provider names
        """
        self._validate_s3_config()
        
        prefix = f"czbillgen/datasets/{dataset_name}/{stage}/"
        response = self.s3_client.list_objects_v2(
            Bucket=self.bucket,
            Prefix=prefix,
            Delimiter="/"
        )
        
        providers = []
        if 'CommonPrefixes' in response:
            for common_prefix in response['CommonPrefixes']:
                # Extract provider from: czbillgen/datasets/{dataset_name}/{stage}/{provider}/
                provider = common_prefix['Prefix'].split('/')[-2]
                providers.append(provider)
        
        return sorted(providers)

    def list_months_with_counts(self, dataset_name: str, stage: str, provider: str) -> List[Tuple[str, int]]:
        """
        List all months for a given dataset+stage+provider with line item counts.
        
        Args:
            dataset_name: Dataset name
            stage: Stage name
            provider: Provider name
            
        Returns:
            List of (month, line_item_count) tuples
        """
        self._validate_s3_config()
        
        prefix = f"czbillgen/datasets/{dataset_name}/{stage}/{provider}/"
        response = self.s3_client.list_objects_v2(
            Bucket=self.bucket,
            Prefix=prefix
        )
        
        months_with_counts = []
        if 'Contents' in response:
            for obj in response['Contents']:
                # Extract month from: czbillgen/datasets/{dataset_name}/{stage}/{provider}/{month}.parquet
                key_parts = obj['Key'].split('/')
                if len(key_parts) >= 5 and key_parts[-1].endswith('.parquet'):
                    month = key_parts[-1].replace('.parquet', '')
                    
                    # Get line item count by reading parquet file - fail-fast on read errors
                    line_count = self._get_line_item_count(obj['Key'])
                    months_with_counts.append((month, line_count))
        
        return sorted(months_with_counts, key=lambda x: x[0])
    
    def list_months_fast(self, dataset_name: str, stage: str, provider: str) -> List[str]:
        """
        List all months for a given dataset+stage+provider without reading files.
        
        Fast version that only checks file existence, used for validation.
        
        Args:
            dataset_name: Dataset name
            stage: Stage name
            provider: Provider name
            
        Returns:
            List of month names (e.g., ['2024-07', '2024-08'])
        """
        self._validate_s3_config()
        
        prefix = f"czbillgen/datasets/{dataset_name}/{stage}/{provider}/"
        response = self.s3_client.list_objects_v2(
            Bucket=self.bucket,
            Prefix=prefix
        )
        
        months = []
        if 'Contents' in response:
            for obj in response['Contents']:
                # Extract month from: czbillgen/datasets/{dataset_name}/{stage}/{provider}/{month}.parquet
                key_parts = obj['Key'].split('/')
                if len(key_parts) >= 5 and key_parts[-1].endswith('.parquet'):
                    month = key_parts[-1].replace('.parquet', '')
                    months.append(month)
        
        return sorted(months)

    def _get_line_item_count(self, key: str) -> int:
        """Get line item count from parquet file."""
        df = self.read_parquet_from_s3(key)
        return len(df)

    def get_month_schema(self, dataset_name: str, stage: str, provider: str, month: str) -> dict:
        """
        Get schema information for a specific month's parquet file.
        
        Args:
            dataset_name: Dataset name
            stage: Stage name
            provider: Provider name
            month: Month (e.g., "2024-07")
            
        Returns:
            Dictionary containing:
            - columns: List of column names
            - dtypes: Dictionary mapping column names to data types
            - row_count: Number of rows in the file
            
        Raises:
            ValueError: If file doesn't exist or can't be read
        """
        self._validate_s3_config()
        
        # Construct the S3 key for the specific month
        key = f"czbillgen/datasets/{dataset_name}/{stage}/{provider}/{month}.parquet"
        
        # Read the parquet file to get schema information
        df = self.read_parquet_from_s3(key)
        
        # Get column names and data types
        columns = df.columns
        dtypes = {col: str(df[col].dtype) for col in columns}
        row_count = len(df)
        
        return {
            "columns": columns,
            "dtypes": dtypes,
            "row_count": row_count
        }
    
    def _validate_basic_dataset_structure(self, data: pl.DataFrame) -> bool:
        """
        Basic dataset validation - just ensure non-empty and valid structure.
        
        Args:
            data: DataFrame to validate
            
        Returns:
            True if basic structure is valid, False otherwise
        """
        if data.height == 0:
            self.logger.error("Dataset is empty")
            return False
        
        if len(data.columns) == 0:
            self.logger.error("Dataset has no columns")
            return False
        
        return True
    
    def _validate_dataset_file(self, file_path: Path) -> bool:
        """
        Basic dataset file validation - just ensure valid parquet format.
        
        Args:
            file_path: Path to dataset file
            
        Returns:
            True if file is valid, False otherwise
        """
        if not file_path.exists() or file_path.stat().st_size == 0:
            self.logger.error(f"Dataset file doesn't exist or is empty: {file_path}")
            return False
        
        # Check valid parquet format and basic structure
        df = pl.read_parquet(file_path)
        return self._validate_basic_dataset_structure(df)
    
    def write_success_marker(self, dataset_name: str, success_content: dict) -> None:
        """
        Write _SUCCESS marker file to indicate complete dataset generation.
        
        Args:
            dataset_name: Name of the dataset
            success_content: Dictionary with success marker information
            
        Raises:
            ValueError: If S3 configuration is invalid
        """
        self._validate_s3_config()
        
        import json
        success_key = f"czbillgen/datasets/{dataset_name}/generate/_SUCCESS"
        
        # Convert success content to JSON
        success_json = json.dumps(success_content, indent=2)
        
        # Upload success marker
        self.s3_client.put_object(
            Bucket=self.bucket,
            Key=success_key,
            Body=success_json.encode('utf-8'),
            ContentType='application/json'
        )
        
        self.logger.info(f"Success marker written: s3://{self.bucket}/{success_key}")

    # Upload-ready dataset navigation methods
    
    def list_upload_organizations(self) -> List[str]:
        """
        List all organizations in piedpiper-custom-billing.
        
        Returns:
            List of organization names
        """
        self._validate_s3_config()
        
        prefix = "piedpiper-custom-billing/"
        response = self.s3_client.list_objects_v2(
            Bucket=self.bucket,
            Prefix=prefix,
            Delimiter='/'
        )
        
        organizations = []
        if 'CommonPrefixes' in response:
            for prefix_obj in response['CommonPrefixes']:
                org_path = prefix_obj['Prefix']
                # Extract org name from: piedpiper-custom-billing/{org}/
                org_name = org_path.split('/')[-2]
                organizations.append(org_name)
        
        return sorted(organizations)
    
    def list_upload_providers(self, organization: str) -> List[str]:
        """
        List all providers for an organization in piedpiper-custom-billing.
        
        Args:
            organization: Organization name
            
        Returns:
            List of provider names
        """
        self._validate_s3_config()
        
        prefix = f"piedpiper-custom-billing/{organization}/"
        response = self.s3_client.list_objects_v2(
            Bucket=self.bucket,
            Prefix=prefix,
            Delimiter='/'
        )
        
        providers = []
        if 'CommonPrefixes' in response:
            for prefix_obj in response['CommonPrefixes']:
                provider_path = prefix_obj['Prefix']
                # Extract provider name from: piedpiper-custom-billing/{org}/{provider}/
                provider_name = provider_path.split('/')[-2]
                providers.append(provider_name)
        
        return sorted(providers)
    
    def list_upload_date_ranges(self, organization: str, provider: str) -> List[str]:
        """
        List all date ranges for an organization and provider in piedpiper-custom-billing.
        
        Args:
            organization: Organization name
            provider: Provider name
            
        Returns:
            List of date ranges (e.g., ['20240701-20240801', '20240801-20240901'])
        """
        self._validate_s3_config()
        
        prefix = f"piedpiper-custom-billing/{organization}/{provider}/"
        response = self.s3_client.list_objects_v2(
            Bucket=self.bucket,
            Prefix=prefix,
            Delimiter='/'
        )
        
        date_ranges = []
        if 'CommonPrefixes' in response:
            for prefix_obj in response['CommonPrefixes']:
                date_range_path = prefix_obj['Prefix']
                # Extract date range from: piedpiper-custom-billing/{org}/{provider}/{date_range}/
                date_range = date_range_path.split('/')[-2]
                date_ranges.append(date_range)
        
        return sorted(date_ranges)
    
    def get_upload_manifest(self, organization: str, provider: str, date_range: str) -> dict:
        """
        Get manifest.json information for a specific date range.
        
        Args:
            organization: Organization name
            provider: Provider name
            date_range: Date range (e.g., '20240701-20240801')
            
        Returns:
            Dictionary with manifest information and row count
        """
        self._validate_s3_config()
        
        import json
        
        # Get manifest.json content
        manifest_key = f"piedpiper-custom-billing/{organization}/{provider}/{date_range}/manifest.json"
        
        try:
            response = self.s3_client.get_object(
                Bucket=self.bucket,
                Key=manifest_key
            )
            manifest_content = response['Body'].read().decode('utf-8')
            manifest_data = json.loads(manifest_content)
        except ClientError as e:
            if e.response['Error']['Code'] == 'NoSuchKey':
                manifest_data = {"error": "manifest.json not found"}
            else:
                raise
        
        # Get data file information
        data_file_info = self._get_upload_data_file_info(organization, provider, date_range)
        
        # Combine manifest and data file information
        result = {
            "manifest": manifest_data,
            "data_file": data_file_info
        }
        
        return result
    
    def _get_upload_data_file_info(self, organization: str, provider: str, date_range: str) -> dict:
        """
        Get information about the data file for a specific date range.
        
        Args:
            organization: Organization name
            provider: Provider name
            date_range: Date range
            
        Returns:
            Dictionary with data file information
        """
        # List files in the date range directory
        prefix = f"piedpiper-custom-billing/{organization}/{provider}/{date_range}/"
        response = self.s3_client.list_objects_v2(
            Bucket=self.bucket,
            Prefix=prefix
        )
        
        if 'Contents' not in response:
            return {"error": "No files found in date range"}
        
        # Find the data file (typically data-export-0001.csv.gz)
        data_file = None
        for obj in response['Contents']:
            if obj['Key'].endswith('.csv.gz') and 'data-export' in obj['Key']:
                data_file = obj
                break
        
        if not data_file:
            return {"error": "No data file found"}
        
        # Get basic file information
        file_info = {
            "file_name": data_file['Key'].split('/')[-1],
            "file_size": data_file['Size'],
            "last_modified": data_file['LastModified'].isoformat()
        }
        
        # Try to get row count and schema from the CSV.gz file
        try:
            import gzip
            import io
            import csv
            
            response = self.s3_client.get_object(
                Bucket=self.bucket,
                Key=data_file['Key']
            )
            
            # Decompress and read CSV
            with gzip.GzipFile(fileobj=io.BytesIO(response['Body'].read())) as gz_file:
                # Read first few lines to get header and sample data
                content = gz_file.read().decode('utf-8')
                lines = content.strip().split('\n')
                
                if lines:
                    # Get column names from header
                    reader = csv.reader([lines[0]])
                    columns = next(reader)
                    
                    # Get row count (excluding header)
                    row_count = len(lines) - 1
                    
                    file_info.update({
                        "columns": columns,
                        "column_count": len(columns),
                        "row_count": row_count
                    })
                else:
                    file_info.update({
                        "columns": [],
                        "column_count": 0,
                        "row_count": 0
                    })
                    
        except Exception as e:
            self.logger.warning(f"Could not read CSV.gz file content: {e}")
            file_info["schema_error"] = str(e)
        
        return file_info

    # Metadata stage-specific operations
    def get_metadata_input_source(self, dataset_name: str, provider: str, month: str) -> str:
        """
        Determine input source (k8s vs generate) for metadata stage.
        
        Args:
            dataset_name: Dataset name
            provider: Cloud provider name
            month: Month in YYYY-MM format
            
        Returns:
            Input source ('k8s' or 'generate')
        """
        # K8s-eligible providers (can have k8s stage data)
        k8s_eligible_providers = ["AWS", "GCP", "Azure"]
        
        # Non-k8s-eligible providers always use generate stage
        if provider not in k8s_eligible_providers:
            return 'generate'
        
        # Check if k8s stage data exists for this provider/month
        k8s_path = f"czbillgen/datasets/{dataset_name}/k8s/{provider}/{month}.parquet"
        
        if self.file_exists(k8s_path):
            self.logger.debug(f"K8s data exists for {provider}/{month}, using k8s stage")
            return 'k8s'
        else:
            self.logger.debug(f"K8s data not found for {provider}/{month}, using generate stage")
            return 'generate'

    def load_metadata_month_data(self, dataset_name: str, provider: str, month: str) -> pl.DataFrame:
        """
        Load input data for metadata stage with dual-source logic.
        
        Args:
            dataset_name: Dataset name
            provider: Cloud provider name
            month: Month in YYYY-MM format
            
        Returns:
            Polars DataFrame with billing data
            
        Raises:
            FileNotFoundError: If input data doesn't exist
            ValueError: If data is corrupted or invalid
        """
        input_source = self.get_metadata_input_source(dataset_name, provider, month)
        
        # Construct S3 path based on input source
        s3_path = f"czbillgen/datasets/{dataset_name}/{input_source}/{provider}/{month}.parquet"
        
        # Read parquet using helper method (handles credential issues)
        df = self.read_parquet_from_s3(s3_path)
        
        # Validate basic schema
        required_columns = ['service', 'resource_id']
        missing_columns = [col for col in required_columns if col not in df.columns]
        
        if missing_columns:
            raise ValueError(f"Missing required columns: {missing_columns}")
        
        if df.height == 0:
            self.logger.warning(f"Empty dataset for {provider}/{month}")
        
        self.logger.info(f"Loaded {df.height:,} rows from {input_source} stage for {provider}/{month}")
        return df

    def save_metadata_month_data(self, dataset_name: str, provider: str, month: str, data: pl.DataFrame) -> None:
        """
        Save processed metadata data to S3.
        
        Args:
            dataset_name: Dataset name
            provider: Cloud provider name
            month: Month in YYYY-MM format
            data: Processed DataFrame with metadata
        """
        # Upload DataFrame directly to S3 metadata stage
        self.upload_dataset_to_stage(dataset_name, "metadata", provider, month, data)
        
        self.logger.info(f"Saved {data.height:,} rows to metadata stage for {provider}/{month}")

    def list_metadata_available_months(self, dataset_name: str) -> List[Tuple[str, str]]:
        """
        List all available provider/month combinations for metadata stage input.
        
        Args:
            dataset_name: Dataset name
            
        Returns:
            List of (provider, month) tuples
        """
        provider_months = []
        
        # Check both k8s and generate stages
        for stage in ['k8s', 'generate']:
            try:
                # Use existing S3DataService methods
                providers = self.list_providers(dataset_name, stage)
                
                for provider in providers:
                    # Get months for this provider in this stage
                    months = self.list_months_fast(dataset_name, stage, provider)
                    
                    for month in months:
                        provider_months.append((provider, month))
                        
            except Exception as e:
                # Continue checking other stages even if one fails
                self.logger.warning(f"Error listing files in {stage} stage: {e}")
                continue
        
        # Remove duplicates and sort
        unique_months = sorted(set(provider_months))
        
        self.logger.info(f"Found {len(unique_months)} provider/month combinations for metadata stage")
        return unique_months

    def create_metadata_success_marker(self, dataset_name: str) -> None:
        """
        DEPRECATED: Create success marker to indicate metadata stage completion.
        
        Use SuccessMarkerService.create_success_marker() instead for standardized format.
        
        Args:
            dataset_name: Dataset name
        """
        success_path = f"czbillgen/datasets/{dataset_name}/metadata/_SUCCESS"
        
        # Create success marker content
        import tempfile
        import os
        
        with tempfile.NamedTemporaryFile(mode='w', delete=False) as tmp_file:
            tmp_file.write(f"Metadata stage completed for dataset: {dataset_name}")
            tmp_path = tmp_file.name
        
        try:
            # Upload to S3
            self.upload_file(tmp_path, success_path)
            
            self.logger.info("Created metadata stage success marker")
            
        finally:
            # Clean up temp file
            if os.path.exists(tmp_path):
                os.unlink(tmp_path)

    def check_metadata_existing_data(self, dataset_name: str, provider: str, month: str) -> bool:
        """
        Check if metadata data already exists for a provider/month.
        
        Args:
            dataset_name: Dataset name
            provider: Cloud provider name
            month: Month in YYYY-MM format
            
        Returns:
            True if data exists, False otherwise
        """
        s3_path = f"czbillgen/datasets/{dataset_name}/metadata/{provider}/{month}.parquet"
        
        return self.file_exists(s3_path)

    # Upload-specific operations
    def load_dataset_for_upload(
        self,
        dataset_name: str,
        stage: str = "generate",
        months: Optional[List[str]] = None,
        providers: Optional[List[str]] = None
    ) -> Dict[str, Dict[str, pl.DataFrame]]:
        """
        Load dataset from S3 with month/provider filtering for upload operations.
        
        Args:
            dataset_name: Name of the dataset to load
            stage: Stage to load from (default: "generate")
            months: List of months to load (format: "YYYY-MM")
            providers: List of providers to load (e.g., ["AWS", "GCP"])
            
        Returns:
            Dict[provider, Dict[month, DataFrame]] - Nested structure for easy processing
            
        Raises:
            FileNotFoundError: If dataset, stage, or requested files not found
            ValueError: If invalid month/provider combinations requested
        """
        # Validate dataset exists
        available_datasets = self.list_datasets()
        if dataset_name not in available_datasets:
            raise FileNotFoundError(f"Dataset '{dataset_name}' not found. Available: {available_datasets}")
        
        # Validate stage exists
        available_stages = self.list_stages(dataset_name)
        if stage not in available_stages:
            raise FileNotFoundError(f"Stage '{stage}' not found in dataset '{dataset_name}'. Available: {available_stages}")
        
        # Get available providers and months
        available_providers = self.list_providers(dataset_name, stage)
        if not available_providers:
            raise FileNotFoundError(f"No providers found for dataset '{dataset_name}' stage '{stage}'")
        
        # Filter providers if specified
        if providers:
            providers_to_load = [p for p in providers if p in available_providers]
            if not providers_to_load:
                raise ValueError(f"None of the requested providers {providers} are available. Available: {available_providers}")
        else:
            providers_to_load = available_providers
        
        # Load data for each provider
        result = {}
        for provider in providers_to_load:
            self.logger.info(f"Loading provider '{provider}' from dataset '{dataset_name}'")
            
            # Get available months for this provider
            available_months_with_counts = self.list_months_with_counts(dataset_name, stage, provider)
            available_months = [month for month, _ in available_months_with_counts]
            
            # Filter months if specified
            if months:
                months_to_load = [m for m in months if m in available_months]
                if not months_to_load:
                    self.logger.warning(f"No requested months {months} found for provider '{provider}'. Available: {available_months}")
                    continue
            else:
                months_to_load = available_months
            
            # Load data for each month
            provider_data = {}
            for month in months_to_load:
                self.logger.info(f"Loading {provider}/{month} from dataset '{dataset_name}'")
                
                try:
                    # Load DataFrame from S3
                    df = self._load_upload_month_data(dataset_name, stage, provider, month)
                    provider_data[month] = df
                    
                    self.logger.info(f"Loaded {len(df)} records for {provider}/{month}")
                    
                except FileNotFoundError as e:
                    # Expected error for missing files - log and continue
                    self.logger.warning(f"File not found for {provider}/{month}: {e}")
                    continue
                except Exception as e:
                    # Unexpected error - log and re-raise to fail fast
                    self.logger.error(f"Unexpected error loading {provider}/{month}: {e}")
                    raise
            
            if provider_data:
                result[provider] = provider_data
                self.logger.info(f"Successfully loaded {len(provider_data)} months for provider '{provider}'")
        
        if not result:
            raise FileNotFoundError(f"No data could be loaded for dataset '{dataset_name}' with the specified filters")
        
        return result

    def _load_upload_month_data(self, dataset_name: str, stage: str, provider: str, month: str) -> pl.DataFrame:
        """Load a specific month's data from S3 for upload operations."""
        # Construct S3 key
        s3_key = f"czbillgen/datasets/{dataset_name}/{stage}/{provider}/{month}.parquet"
        
        # Load from S3 using helper method (works around Polars S3 issues)
        df = self.read_parquet_from_s3(s3_key)
        
        # Basic validation
        if len(df) == 0:
            raise ValueError(f"Empty dataset found at s3://{self.bucket}/{s3_key}")
        
        return df

    def get_upload_dataset_summary(self, dataset_name: str, stage: str = "generate") -> Dict[str, Dict[str, int]]:
        """
        Get summary of dataset contents for upload operations without loading data.
        
        Args:
            dataset_name: Name of the dataset
            stage: Stage to summarize (default: "generate")
            
        Returns:
            Dict[provider, Dict[month, line_count]] - Summary of available data
        """
        # Validate dataset and stage exist
        available_datasets = self.list_datasets()
        if dataset_name not in available_datasets:
            raise FileNotFoundError(f"Dataset '{dataset_name}' not found. Available: {available_datasets}")
        
        available_stages = self.list_stages(dataset_name)
        if stage not in available_stages:
            raise FileNotFoundError(f"Stage '{stage}' not found in dataset '{dataset_name}'. Available: {available_stages}")
        
        # Get summary for each provider
        summary = {}
        available_providers = self.list_providers(dataset_name, stage)
        
        for provider in available_providers:
            months_with_counts = self.list_months_with_counts(dataset_name, stage, provider)
            provider_summary = {month: count for month, count in months_with_counts}
            summary[provider] = provider_summary
        
        return summary

    def get_upload_dataset_summary_fast(self, dataset_name: str, stage: str = "generate") -> Dict[str, List[str]]:
        """
        Get fast summary of dataset contents for upload operations without loading data or counting rows.
        
        Args:
            dataset_name: Name of the dataset
            stage: Stage to summarize (default: "generate")
            
        Returns:
            Dict[provider, List[month]] - Summary of available files without row counts
        """
        # Validate dataset and stage exist
        available_datasets = self.list_datasets()
        if dataset_name not in available_datasets:
            raise FileNotFoundError(f"Dataset '{dataset_name}' not found. Available: {available_datasets}")
        
        available_stages = self.list_stages(dataset_name)
        if stage not in available_stages:
            raise FileNotFoundError(f"Stage '{stage}' not found in dataset '{dataset_name}'. Available: {available_stages}")
        
        # Get summary for each provider (fast version)
        summary = {}
        available_providers = self.list_providers(dataset_name, stage)
        
        for provider in available_providers:
            # Use fast method that only checks file existence
            months = self.list_months_fast(dataset_name, stage, provider)
            summary[provider] = months
        
        return summary

    def validate_upload_dataset_access(self, dataset_name: str, stage: str = "generate") -> bool:
        """
        Validate that we can access the dataset for upload operations without loading it.
        
        Args:
            dataset_name: Name of the dataset
            stage: Stage to validate (default: "generate")
            
        Returns:
            True if dataset is accessible, False otherwise
        """
        # Check if dataset exists
        available_datasets = self.list_datasets()
        if dataset_name not in available_datasets:
            self.logger.error(f"Dataset '{dataset_name}' not found")
            return False
        
        # Check if stage exists
        available_stages = self.list_stages(dataset_name)
        if stage not in available_stages:
            self.logger.error(f"Stage '{stage}' not found in dataset '{dataset_name}'")
            return False
        
        # Check if there are any providers
        available_providers = self.list_providers(dataset_name, stage)
        if not available_providers:
            self.logger.error(f"No providers found for dataset '{dataset_name}' stage '{stage}'")
            return False
        
        # Check if there are any months for at least one provider
        has_data = False
        for provider in available_providers:
            months_with_counts = self.list_months_with_counts(dataset_name, stage, provider)
            if months_with_counts:
                has_data = True
                break
        
        if not has_data:
            self.logger.error(f"No data found for dataset '{dataset_name}' stage '{stage}'")
            return False
        
        return True
    
    def list_drops_in_date_range(self, organization: str, provider: str, date_range: str) -> List[Dict[str, Any]]:
        """
        List all drops in a date range with aggregated file information.
        
        Args:
            organization: Organization name
            provider: Provider name  
            date_range: Date range (e.g., '20240701-20240801')
            
        Returns:
            List of drop information dictionaries
        """
        self._validate_s3_config()
        
        import json
        
        # Get manifest to identify active drop
        manifest_key = f"piedpiper-custom-billing/{organization}/{provider}/{date_range}/manifest.json"
        active_drop_id = None
        
        try:
            response = self.s3_client.get_object(Bucket=self.bucket, Key=manifest_key)
            manifest_data = json.loads(response['Body'].read().decode('utf-8'))
            active_drop_id = manifest_data.get("current_drop_id")
        except:
            pass  # Continue without active drop info
        
        # List all directories (potential drops) in the date range
        prefix = f"piedpiper-custom-billing/{organization}/{provider}/{date_range}/"
        response = self.s3_client.list_objects_v2(
            Bucket=self.bucket,
            Prefix=prefix,
            Delimiter='/'
        )
        
        drops = []
        if 'CommonPrefixes' in response:
            for prefix_info in response['CommonPrefixes']:
                drop_path = prefix_info['Prefix']
                # Extract drop ID from path
                drop_id = drop_path.rstrip('/').split('/')[-1]
                
                # Skip if it's not a drop directory (e.g., manifest.json is at root level)
                if not drop_id or drop_id in ['manifest.json']:
                    continue
                
                # Get aggregated file info for this drop
                drop_info = self._get_drop_file_info(organization, provider, date_range, drop_id)
                drop_info.update({
                    "drop_id": drop_id,
                    "is_active": drop_id == active_drop_id
                })
                
                drops.append(drop_info)
        
        # Sort drops by drop_id (newest first, assuming timestamp format)
        drops.sort(key=lambda x: x['drop_id'], reverse=True)
        
        return drops
    
    def _get_drop_file_info(self, organization: str, provider: str, date_range: str, drop_id: str) -> Dict[str, Any]:
        """
        Get aggregated file information for a specific drop.
        
        Args:
            organization: Organization name
            provider: Provider name
            date_range: Date range
            drop_id: Drop ID
            
        Returns:
            Dictionary with aggregated file information
        """
        try:
            # List all files in the drop directory
            prefix = f"piedpiper-custom-billing/{organization}/{provider}/{date_range}/{drop_id}/"
            response = self.s3_client.list_objects_v2(
                Bucket=self.bucket,
                Prefix=prefix
            )
            
            files = []
            total_size = 0
            total_rows = 0
            
            if 'Contents' in response:
                for obj in response['Contents']:
                    if obj['Key'].endswith('.csv.gz') and 'data-export' in obj['Key']:
                        filename = obj['Key'].split('/')[-1]
                        file_size = obj['Size']
                        last_modified = obj['LastModified']
                        
                        # Get row count (optimized - only read if needed)
                        row_count = self._get_file_row_count_optimized(obj['Key'])
                        
                        files.append({
                            "filename": filename,
                            "size": file_size,
                            "rows": row_count,
                            "last_modified": last_modified
                        })
                        
                        total_size += file_size
                        total_rows += row_count
            
            return {
                "files": files,
                "file_count": len(files),
                "total_size": total_size,
                "total_rows": total_rows,
                "upload_date": files[0]["last_modified"] if files else None
            }
            
        except Exception as e:
            self.logger.error(f"Failed to get drop file info for {drop_id}: {e}")
            return {
                "files": [],
                "file_count": 0,
                "total_size": 0,
                "total_rows": 0,
                "error": str(e)
            }
    
    def _get_file_row_count_optimized(self, s3_key: str) -> int:
        """
        Get row count from gzipped CSV file with optimization.
        For performance, only count lines without parsing full CSV.
        """
        try:
            response = self.s3_client.get_object(Bucket=self.bucket, Key=s3_key)
            
            import gzip
            import io
            
            # Count lines in gzipped file
            with gzip.open(io.BytesIO(response['Body'].read()), 'rt', encoding='utf-8') as gz_file:
                # Count lines efficiently
                row_count = sum(1 for _ in gz_file) - 1  # Subtract 1 for header
                return max(0, row_count)  # Ensure non-negative
                
        except Exception as e:
            self.logger.warning(f"Failed to get row count for {s3_key}: {e}")
            return 0
    
    def get_parquet_row_count(self, s3_key: str) -> int:
        """
        Get row count from Parquet file without loading full dataset.
        Uses Polars metadata scanning for efficiency.
        """
        try:
            # Use the same credential handling as scan_parquet_from_s3
            lazy_df = self.scan_parquet_from_s3(s3_key)
            return lazy_df.select(pl.count()).collect().item()
        except Exception as e:
            self.logger.warning(f"Failed to get parquet row count for {s3_key}: {e}")
            return 0
    
    def scan_parquet_from_s3(self, s3_key: str) -> pl.LazyFrame:
        """
        Create a lazy scan of Parquet file from S3.
        Returns LazyFrame that can be processed with Polars lazy operations.
        """
        # Get AWS credentials from boto3 session to pass to Polars
        if self.profile:
            session = boto3.Session(profile_name=self.profile)
        else:
            session = boto3.Session()
        
        # Get credentials from the session
        credentials = session.get_credentials()
        if credentials:
            s3_options = {
                "region": self.region,
                "access_key_id": credentials.access_key,
                "secret_access_key": credentials.secret_key
            }
            # Add session token if available (needed for temporary credentials)
            if credentials.token:
                s3_options["session_token"] = credentials.token
        else:
            # Fallback to just region if no credentials found
            s3_options = {
                "region": self.region
            }
        
        s3_path = f"s3://{self.bucket}/{s3_key}"
        return pl.scan_parquet(s3_path, storage_options=s3_options)
    
    def get_drop_details(self, organization: str, provider: str, date_range: str, drop_id: str) -> Dict[str, Any]:
        """
        Get detailed information for a specific drop including file details and column schema.
        
        Args:
            organization: Organization name
            provider: Provider name
            date_range: Date range
            drop_id: Drop ID
            
        Returns:
            Dictionary with detailed drop information
        """
        self._validate_s3_config()
        
        try:
            # Get basic drop info
            drop_info = self._get_drop_file_info(organization, provider, date_range, drop_id)
            
            # Check if this is the active drop
            active_drop_id = self._get_active_drop_id(organization, provider, date_range)
            drop_info["is_active"] = (drop_id == active_drop_id)
            drop_info["drop_id"] = drop_id
            
            # Transform file details for CLI display
            if drop_info["files"]:
                # Update file details with proper field names for CLI
                for file_detail in drop_info["files"]:
                    file_detail["size_bytes"] = file_detail.pop("size")
                    file_detail["last_modified"] = file_detail["last_modified"].isoformat() if hasattr(file_detail["last_modified"], 'isoformat') else str(file_detail["last_modified"])
                
                # Get column schema from first file
                first_file = drop_info["files"][0]
                s3_key = f"piedpiper-custom-billing/{organization}/{provider}/{date_range}/{drop_id}/{first_file['filename']}"
                
                try:
                    # Get column names from first file
                    response = self.s3_client.get_object(Bucket=self.bucket, Key=s3_key)
                    
                    import gzip
                    import io
                    import csv
                    
                    with gzip.open(io.BytesIO(response['Body'].read()), 'rt', encoding='utf-8') as gz_file:
                        # Read just the header
                        csv_reader = csv.reader(gz_file)
                        columns = next(csv_reader, [])
                        drop_info["columns"] = columns
                        drop_info["column_count"] = len(columns)
                        
                except Exception as e:
                    self.logger.warning(f"Failed to get column schema: {e}")
                    drop_info["columns"] = []
                    drop_info["column_count"] = 0
                    drop_info["schema_error"] = str(e)
            else:
                drop_info["columns"] = []
                drop_info["column_count"] = 0
            
            return drop_info
            
        except Exception as e:
            self.logger.error(f"Failed to get drop details for {drop_id}: {e}")
            return {"error": str(e)}
    
    def list_drops_in_date_range(self, organization: str, provider: str, date_range: str) -> List[Dict[str, Any]]:
        """
        List all drops in a date range with aggregated file information.
        
        Args:
            organization: Organization name
            provider: Provider name
            date_range: Date range
            
        Returns:
            List of dictionaries with drop information including active drop indication
        """
        self._validate_s3_config()
        
        try:
            # Get active drop ID from manifest
            active_drop_id = self._get_active_drop_id(organization, provider, date_range)
            
            # List all drops in the date range
            prefix = f"piedpiper-custom-billing/{organization}/{provider}/{date_range}/"
            response = self.s3_client.list_objects_v2(
                Bucket=self.bucket,
                Prefix=prefix,
                Delimiter='/'
            )
            
            drops = []
            if 'CommonPrefixes' in response:
                for prefix_obj in response['CommonPrefixes']:
                    drop_path = prefix_obj['Prefix']
                    # Extract drop ID from path
                    drop_id = drop_path.split('/')[-2]
                    
                    # Skip non-drop folders (like manifest.json level)
                    if drop_id and drop_id != date_range and not drop_id.endswith('.json'):
                        # Get aggregated file info for this drop
                        drop_info = self._get_drop_file_info(organization, provider, date_range, drop_id)
                        
                        drops.append({
                            "drop_id": drop_id,
                            "file_count": drop_info["file_count"],
                            "total_rows": drop_info["total_rows"],
                            "active_drop_id": active_drop_id  # Include for comparison in CLI
                        })
            
            # Sort by drop ID (reverse chronological since drop IDs are timestamps)
            drops.sort(key=lambda x: x["drop_id"], reverse=True)
            
            return drops
            
        except Exception as e:
            self.logger.error(f"Failed to list drops for {organization}/{provider}/{date_range}: {e}")
            raise ValueError(f"Unable to list drops for {organization}/{provider}/{date_range}: {e}")
    
    def _get_active_drop_id(self, organization: str, provider: str, date_range: str) -> Optional[str]:
        """
        Get the active drop ID from manifest.json.
        
        Args:
            organization: Organization name
            provider: Provider name
            date_range: Date range
            
        Returns:
            Active drop ID or None if not found
        """
        try:
            manifest_key = f"piedpiper-custom-billing/{organization}/{provider}/{date_range}/manifest.json"
            
            response = self.s3_client.get_object(
                Bucket=self.bucket,
                Key=manifest_key
            )
            
            manifest_content = response['Body'].read().decode('utf-8')
            manifest_data = json.loads(manifest_content)
            
            return manifest_data.get('current_drop_id')
            
        except Exception as e:
            self.logger.warning(f"Failed to get active drop ID: {e}")
            return None

    # Version Management Methods

    def list_version_files(self) -> List[Tuple[str, int, datetime]]:
        """
        List all version files in S3 with metadata.

        Returns:
            List of (filename, size_bytes, modified_time) tuples
        """
        try:
            response = self.s3_client.list_objects_v2(
                Bucket=self.bucket,
                Prefix=f"{self.reference_path}/versions/"
            )

            if 'Contents' not in response:
                return []

            version_files = []
            for obj in response['Contents']:
                key = obj['Key']
                filename = key.split('/')[-1]  # Get just the filename

                # Skip current.parquet and non-parquet files
                if filename == 'current.parquet' or not filename.endswith('.parquet'):
                    continue

                version_files.append((
                    filename,
                    obj['Size'],
                    obj['LastModified']
                ))

            return version_files

        except Exception as e:
            self.logger.error(f"Failed to list version files: {e}")
            return []

    def version_file_exists(self, version_id: str) -> bool:
        """
        Check if a specific version file exists in S3.

        Args:
            version_id: Version identifier

        Returns:
            True if version file exists, False otherwise
        """
        try:
            version_key = f"{self.reference_path}/versions/{version_id}.parquet"
            self.s3_client.head_object(Bucket=self.bucket, Key=version_key)
            return True
        except ClientError as e:
            if e.response['Error']['Code'] == '404':
                return False
            raise
        except Exception as e:
            self.logger.error(f"Error checking version file {version_id}: {e}")
            return False

    def get_current_file_info(self) -> Optional[Tuple[int, datetime]]:
        """
        Get metadata for the current reference data file.

        Returns:
            (file_size_bytes, modified_time) or None if not found
        """
        try:
            current_key = f"{self.reference_path}/current.parquet"
            response = self.s3_client.head_object(Bucket=self.bucket, Key=current_key)
            return response['ContentLength'], response['LastModified']
        except ClientError as e:
            if e.response['Error']['Code'] == '404':
                return None
            raise
        except Exception as e:
            self.logger.error(f"Error getting current file info: {e}")
            return None

    def create_version_backup(self, version_id: str) -> bool:
        """
        Create a versioned backup of the current reference data.

        Args:
            version_id: Version identifier for the backup

        Returns:
            True if backup created successfully, False otherwise
        """
        try:
            current_key = f"{self.reference_path}/current.parquet"
            version_key = f"{self.reference_path}/versions/{version_id}.parquet"

            # Copy current to version
            copy_source = {'Bucket': self.bucket, 'Key': current_key}
            self.s3_client.copy_object(
                CopySource=copy_source,
                Bucket=self.bucket,
                Key=version_key
            )

            self.logger.info(f"Created version backup: {version_id}")
            return True

        except Exception as e:
            self.logger.error(f"Failed to create version backup {version_id}: {e}")
            return False

    def restore_version(self, version_id: str) -> bool:
        """
        Restore a specific version as the current reference data.

        Args:
            version_id: Version to restore

        Returns:
            True if restore successful, False otherwise
        """
        try:
            version_key = f"{self.reference_path}/versions/{version_id}.parquet"
            current_key = f"{self.reference_path}/current.parquet"

            # Copy version to current
            copy_source = {'Bucket': self.bucket, 'Key': version_key}
            self.s3_client.copy_object(
                CopySource=copy_source,
                Bucket=self.bucket,
                Key=current_key
            )

            # Update version metadata
            self._update_version_metadata({'current_version': version_id})

            self.logger.info(f"Restored version: {version_id}")
            return True

        except Exception as e:
            self.logger.error(f"Failed to restore version {version_id}: {e}")
            return False

    def get_version_metadata(self) -> Optional[Dict[str, Any]]:
        """
        Get version metadata from S3.

        Returns:
            Version metadata dict or None if not found
        """
        try:
            metadata_key = f"{self.reference_path}/version_metadata.json"
            response = self.s3_client.get_object(Bucket=self.bucket, Key=metadata_key)
            content = response['Body'].read().decode('utf-8')
            return json.loads(content)
        except ClientError as e:
            if e.response['Error']['Code'] == '404':
                return None
            raise
        except Exception as e:
            self.logger.error(f"Error getting version metadata: {e}")
            return None

    def _update_version_metadata(self, metadata: Dict[str, Any]) -> None:
        """
        Update version metadata in S3.

        Args:
            metadata: Metadata dict to store
        """
        try:
            metadata_key = f"{self.reference_path}/version_metadata.json"

            # Merge with existing metadata
            existing = self.get_version_metadata() or {}
            existing.update(metadata)
            existing['updated_at'] = datetime.utcnow().isoformat()

            # Upload updated metadata
            self.s3_client.put_object(
                Bucket=self.bucket,
                Key=metadata_key,
                Body=json.dumps(existing, indent=2),
                ContentType='application/json'
            )

        except Exception as e:
            self.logger.error(f"Failed to update version metadata: {e}")
            # Don't raise - metadata update failure shouldn't block operations

    # Methods needed by VersionManager for simplified local-first architecture
    def get_config_value(self, key: str) -> str:
        """Get configuration value needed by VersionManager."""
        if key == "reference_data_bucket":
            return self.bucket
        elif key == "reference_data_prefix":
            return self.reference_path
        else:
            raise ValueError(f"Unknown config key: {key}")

    def get_local_cache_path(self) -> Path:
        """Get path to local cache file."""
        return self._get_cache_file_path("reference.parquet")

    def download_reference_data_to_path(self, local_path: Path) -> None:
        """Download current reference data to specified local path."""
        current_key = self._get_s3_reference_key("current.parquet")

        try:
            self.s3_client.download_file(
                Bucket=self.bucket,
                Key=current_key,
                Filename=str(local_path)
            )
            self.logger.debug(f"Downloaded reference data to {local_path}")
        except ClientError as e:
            if e.response['Error']['Code'] == '404':
                raise FileNotFoundError(f"No current reference data found in S3")
            raise

    # Dataset metadata operations
    def read_text_file(self, s3_key: str) -> str:
        """Read text content from S3."""
        self._validate_s3_config()
        try:
            response = self.s3_client.get_object(Bucket=self.bucket, Key=s3_key)
            return response['Body'].read().decode('utf-8')
        except ClientError as e:
            if e.response['Error']['Code'] == '404':
                raise FileNotFoundError(f"File not found in S3: {s3_key}")
            raise

    def upload_text_content(self, s3_key: str, content: str) -> None:
        """Upload text content to S3."""
        self._validate_s3_config()
        self.s3_client.put_object(
            Bucket=self.bucket,
            Key=s3_key,
            Body=content.encode('utf-8'),
            ContentType='text/plain'
        )

    def list_folders(self, prefix: str) -> List[str]:
        """List folders (directories) under a given prefix."""
        self._validate_s3_config()
        try:
            response = self.s3_client.list_objects_v2(
                Bucket=self.bucket,
                Prefix=prefix,
                Delimiter='/'
            )

            folders = []
            if 'CommonPrefixes' in response:
                for obj in response['CommonPrefixes']:
                    folders.append(obj['Prefix'])

            return folders
        except Exception as e:
            self.logger.error(f"Failed to list folders under {prefix}: {e}")
            return []

    def list_files(self, prefix: str, max_files: Optional[int] = None) -> List[str]:
        """List files under a given prefix."""
        self._validate_s3_config()
        try:
            kwargs = {
                'Bucket': self.bucket,
                'Prefix': prefix
            }
            if max_files:
                kwargs['MaxKeys'] = max_files

            response = self.s3_client.list_objects_v2(**kwargs)

            if 'Contents' not in response:
                return []

            files = [obj['Key'] for obj in response['Contents']]
            return files
        except Exception as e:
            self.logger.error(f"Failed to list files under {prefix}: {e}")
            return []