"""
S3 upload orchestrator for multi-org dataset publishing.

Coordinates uploading a single generated dataset to multiple organizations
with proper CBF processing, S3 formatting, and error handling.
"""

from typing import Dict, Any, List
from pathlib import Path
import polars as pl
from rich.console import Console
from rich.table import Table

from .config import UploadConfig, Organization, S3Config
from .uploader import S3Uploader, S3UploadError
from .formatter import S3Formatter
from .folder_manager import S3FolderManager
from .stage_discovery import StageDiscoveryService, StageDiscoveryError, PartialDataError, NoDataError
from ..core.data_processor import DataProcessor
from ..core.connection_manager import ConnectionManager
from ..core.connection_recreator import ConnectionRecreator
from ...services.s3_data_service import S3DataService


class S3UploadOrchestrator:
    """
    S3 upload orchestrator for generated and ad-hoc datasets.
    
    Coordinates the complete upload workflow:
    1. Load and process datasets (generated or ad-hoc) with CBF mapping
    2. Split by month and filter by providers per organization
    3. Upload to S3 with drop IDs and manifest management
    4. Handle multi-org coordination with continue-on-failure
    """
    
    def __init__(self, config: UploadConfig):
        """Initialize orchestrator with configuration."""
        self.config = config
        self.s3_data_service = S3DataService()
        self.s3_formatter = S3Formatter()
        self.console = Console()
        
        # Initialize S3 uploader with environment configuration
        # Determine upload type from config
        upload_type = "generated" if config.is_generated_upload() else "adhoc"
        self.s3_config = S3Config.from_environment(upload_type)
        self.s3_uploader = S3Uploader(self.s3_config)
        
        # Initialize connection manager
        self.connection_manager = ConnectionManager()
        
        # Initialize S3 dataset loader  
        self.dataset_loader = S3DataService()
        
        # Initialize recreate connections services
        self.folder_manager = S3FolderManager(self.s3_config.bucket, self.s3_config.region)
        self.connection_recreator = ConnectionRecreator(self.s3_config.bucket, self.s3_config.region)
        
        # Initialize stage discovery service
        self.stage_discovery = StageDiscoveryService()
        
        # Initialize data processor only for non-dry-run (will be set later)
        self.data_processor = None
    
    def upload_dataset_to_orgs(self, dry_run: bool = False) -> Dict[str, Any]:
        """
        Upload generated dataset to all configured organizations.
        
        Uses unified validation and streaming processing for both dry-run and actual uploads.
        Implements continue-on-failure pattern where individual organization
        failures don't stop processing of other organizations.
        """
        if dry_run:
            print(f"🧪 Validating dataset: {self.config.dataset_name}")
        else:
            print(f"🚀 Starting S3 upload for dataset: {self.config.dataset_name}")
            print(f"🏢 Organizations: {len(self.config.organizations)}")
            # Initialize data processor only for real uploads (loads CBF mapping)
            self.data_processor = DataProcessor()
        
        # Phase 1: Perform unified validation upfront
        print(f"\n🔍 Phase 1: Performing unified validation...")
        validation_result = self._perform_unified_validation(dry_run)
        
        if validation_result["status"] != "success":
            return {
                "status": "failed",
                "error": validation_result["error"],
                "dataset": self.config.dataset_name,
                "organizations_processed": 0,
                "organizations_successful": 0,
                "results": {},
                "validation_step": validation_result.get("validation_step", "unknown")
            }
        
        # Extract validation results
        org_validation_results = validation_result["org_validation_results"]
        processing_plan = validation_result["processing_plan"]
        
        print(f"✅ Validation complete. Processing plan: {processing_plan['total_files']} files across {processing_plan['total_organizations']} organizations")
        
        # Phase 2: Process each organization (continue-on-failure)
        print(f"\n🚀 Phase 2: Processing organizations...")
        results = {}
        successful_orgs = 0
        total_orgs = len(self.config.organizations)
        
        for i, org_plan in enumerate(processing_plan["organizations"], 1):
            org_name = org_plan["name"]
            org = next(org for org in self.config.organizations if org.name == org_name)
            
            if dry_run:
                print(f"\n🏢 Validating organization {i}/{total_orgs}: {org_name}")
            else:
                print(f"\n🏢 Processing organization {i}/{total_orgs}: {org_name}")
            
            # Skip organizations that failed validation
            if org_plan["status"] != "ready":
                results[org_name] = {
                    "status": "failed",
                    "error": org_plan.get("error", "Validation failed")
                }
                print(f"❌ {org_name}: {org_plan.get('error', 'Validation failed')}")
                continue
            
            # Handle missing connections for actual uploads
            if not dry_run:
                try:
                    providers = org.providers  # Now a simple list
                    
                    if providers:
                        # Ensure connections exist for all providers
                        from ..api.connection_manager import get_org_api_key, ConnectionManager
                        
                        org_api_key = get_org_api_key(org.name)
                        conn_manager = ConnectionManager(org_api_key)
                        
                        # Prepare S3 config for connection creation
                        s3_config_dict = {
                            "bucket_name": self.s3_config.bucket,
                            "root_folder": self.s3_config.root_folder,
                            "org_folder": org.name  # Use organization name as folder
                        }
                        
                        # Ensure all required connections exist
                        print(f"   🔍 Ensuring connections for {org.name}...")
                        conn_results = conn_manager.ensure_all_connections(providers, s3_config_dict)
                        print(f"   ✅ Connection check complete for {len(providers)} providers")
                        
                except Exception as e:
                    error_msg = f"Connection management failed for {org.name}: {str(e)}"
                    results[org.name] = {
                        "status": "failed",
                        "error": error_msg
                    }
                    print(f"❌ {org.name}: {error_msg}")
                    continue
            
            # Process organization using appropriate method
            try:
                if dry_run:
                    org_result = self._process_dry_run_organization(org, org_plan)
                else:
                    org_result = self._stream_process_organization(org, org_plan)
                
                results[org_name] = org_result
                
                if org_result["status"] == "success":
                    successful_orgs += 1
                    if dry_run:
                        print(f"✅ {org_name}: Validation completed successfully")
                    else:
                        print(f"✅ {org_name}: Upload completed successfully")
                else:
                    if dry_run:
                        print(f"❌ {org_name}: Validation failed - {org_result.get('error', 'Unknown error')}")
                    else:
                        print(f"❌ {org_name}: Upload failed - {org_result.get('error', 'Unknown error')}")
                    
            except Exception as e:
                error_msg = f"Unexpected error during processing: {str(e)}"
                results[org_name] = {
                    "status": "failed",
                    "error": error_msg
                }
                print(f"❌ {org_name}: {error_msg}")
        
        # Determine overall status
        overall_status = "success" if successful_orgs == total_orgs else "partial" if successful_orgs > 0 else "failed"
        
        # Show concise summary
        status_icon = "✅" if overall_status == "success" else "⚠️" if overall_status == "partial" else "❌"
        if dry_run:
            mode_text = "VALIDATION"
            status_word = "validated" if overall_status == "success" else "failed validation"
        else:
            mode_text = "UPLOAD"
            status_word = overall_status
        print(f"\n{status_icon} {mode_text} SUMMARY: {successful_orgs}/{total_orgs} organizations {status_word}")
        
        # Show detailed results only if there are failures or in verbose mode
        if overall_status != "success" or dry_run:
            for org_name, org_result in results.items():
                result_icon = "✅" if org_result["status"] == "success" else "❌"
                if org_result["status"] == "success":
                    rows = org_result.get('total_rows', 0)
                    files = org_result.get('files_uploaded', 0)
                    if rows == "N/A (dry-run)":
                        print(f"   {result_icon} {org_name}: {files} files")
                    else:
                        print(f"   {result_icon} {org_name}: {files} files, {rows:,} rows")
                else:
                    print(f"   {result_icon} {org_name}: {org_result.get('error', 'Unknown error')}")
        
        print(f"🎯 Dataset: {self.config.dataset_name}")
        
        return {
            "status": overall_status,
            "dataset": self.config.dataset_name,
            "organizations_processed": total_orgs,
            "organizations_successful": successful_orgs,
            "results": results
        }
    
    def upload_adhoc_to_orgs(self, dry_run: bool = False) -> Dict[str, Any]:
        """
        Upload ad-hoc CSV dataset to all configured organizations.
        
        Implements continue-on-failure pattern where individual organization
        failures don't stop processing of other organizations.
        """
        print(f"🚀 Starting ad-hoc upload for CSV: {self.config.adhoc.csv_path}")
        print(f"🏢 Organizations: {len(self.config.organizations)}")
        
        if dry_run:
            print("🧪 DRY RUN MODE - No files will be uploaded")
        
        # Validate S3 bucket access
        try:
            self.s3_uploader.validate_bucket_access()
        except S3UploadError as e:
            return {
                "status": "failed",
                "error": f"S3 bucket validation failed: {str(e)}",
                "csv_path": self.config.adhoc.csv_path,
                "organizations_processed": 0,
                "organizations_successful": 0,
                "results": {}
            }
        
        # Load and process ad-hoc dataset
        try:
            print(f"[•] Loading ad-hoc CSV: {self.config.adhoc.csv_path}")
            csv_path = Path(self.config.adhoc.csv_path)
            raw_df = self.dataset_loader.load_adhoc_dataset(csv_path)
            
            # Load user mapping if provided
            user_mapping = []
            if self.config.adhoc.mapping_path:
                print(f"[•] Loading user mapping: {self.config.adhoc.mapping_path}")
                mapping_path = Path(self.config.adhoc.mapping_path)
                user_mapping = self.dataset_loader.load_user_mapping(mapping_path)
            
            # Process ad-hoc dataset with CBF mapping
            processed_df = self.data_processor.process_adhoc_dataset(raw_df, user_mapping)
            
            print(f"[✓] Processed {len(processed_df)} rows from ad-hoc CSV")
            
        except Exception as e:
            return {
                "status": "failed",
                "error": f"Failed to load/process ad-hoc dataset: {str(e)}",
                "csv_path": self.config.adhoc.csv_path,
                "organizations_processed": 0,
                "organizations_successful": 0,
                "results": {}
            }
        
        # Upload to each organization
        results = {}
        successful_orgs = 0
        
        for org in self.config.organizations:
            print(f"\n📤 Processing organization: {org.name}")
            
            # Auto-create missing connections for this organization
            try:
                print(f"   🔍 Checking connections for {org.name}...")
                providers = list(org.providers.keys())
                
                if providers:
                    # Ensure connections exist for all providers
                    from ..api.connection_manager import get_org_api_key, ConnectionManager
                    
                    org_api_key = get_org_api_key(org.name)
                    conn_manager = ConnectionManager(org_api_key)
                    
                    # Prepare S3 config for connection creation
                    s3_config_dict = {
                        "bucket_name": self.s3_config.bucket,
                        "root_folder": self.s3_config.root_folder,
                        "org_folder": org.name  # Use organization name as folder
                    }
                    
                    if not dry_run:
                        # Ensure all required connections exist
                        conn_results = conn_manager.ensure_all_connections(providers, s3_config_dict)
                        print(f"   ✅ Connection check complete for {len(providers)} providers")
                    else:
                        # In dry run, just validate existing connections
                        conn_status = conn_manager.validate_connections_exist(providers)
                        missing_providers = [p for p, exists in conn_status.items() if not exists]
                        if missing_providers:
                            print(f"   ⚠️  DRY RUN: Would create missing connections for: {missing_providers}")
                        else:
                            print(f"   ✅ All connections exist for providers: {providers}")
                            
            except Exception as e:
                error_msg = f"Connection management failed for {org.name}: {str(e)}"
                results[org.name] = {
                    "status": "failed",
                    "error": error_msg
                }
                print(f"❌ {org.name}: {error_msg}")
                continue
            
            try:
                result = self._upload_adhoc_to_organization(org, processed_df, dry_run)
                results[org.name] = result
                
                if result["status"] == "success":
                    successful_orgs += 1
                    print(f"✅ {org.name}: Upload successful")
                else:
                    print(f"❌ {org.name}: Upload failed - {result.get('error', 'Unknown error')}")
                    
            except Exception as e:
                results[org.name] = {
                    "status": "failed",
                    "error": f"Unexpected error: {str(e)}"
                }
                print(f"❌ {org.name}: Unexpected error - {str(e)}")
        
        # Determine overall status
        total_orgs = len(self.config.organizations)
        if successful_orgs == total_orgs:
            status = "success"
        elif successful_orgs > 0:
            status = "partial"
        else:
            status = "failed"
        
        return {
            "status": status,
            "csv_path": self.config.adhoc.csv_path,
            "organizations_processed": total_orgs,
            "organizations_successful": successful_orgs,
            "results": results
        }
    
    
    def _process_dry_run_organization(
        self, 
        org: Organization,
        org_plan: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Process dry-run validation for an organization using the processing plan.
        
        Uses the pre-validated processing plan to show what would be uploaded.
        """
        estimated_files = 0
        table_data = []
        
        # Process each provider from the plan
        for provider_plan in org_plan["providers"]:
            provider = provider_plan["name"]
            months = provider_plan["months"]
            
            estimated_files += len(months)
            
            # Format months for table display (10 per line)
            month_lines = []
            for i in range(0, len(months), 10):
                chunk = months[i:i+10]
                month_lines.append(", ".join(chunk))
            
            table_data.append({
                "provider": provider,
                "month_lines": month_lines
            })
        
        # Create and display Rich table
        table = Table(show_header=True, header_style="bold blue")
        table.add_column("Provider", style="cyan", width=15)
        table.add_column("Months", style="white", width=100)
        
        for provider_data in table_data:
            provider = provider_data["provider"]
            month_lines = provider_data["month_lines"]
            
            # Add first row with provider name
            if month_lines:
                table.add_row(provider, month_lines[0])
                
                # Add continuation rows for additional month lines
                for month_line in month_lines[1:]:
                    table.add_row("", month_line)
                
                # Add empty separator row between providers (except for the last one)
                if provider != table_data[-1]["provider"]:
                    table.add_row("", "")
        
        print("   ")
        self.console.print(table)
        print(f"   🧪 DRY RUN: Would upload {estimated_files} files")
        
        return {
            "status": "success",
            "organization": org.name,
            "drop_id": "N/A (dry-run)",
            "files_uploaded": estimated_files,
            "total_rows": "N/A (dry-run)",
            "providers": [p["name"] for p in org_plan["providers"]],
            "dry_run": True
        }
    
    def _upload_adhoc_to_organization(
        self, 
        org: Organization,
        df: pl.DataFrame,
        dry_run: bool
    ) -> Dict[str, Any]:
        """Upload ad-hoc dataset to a single organization."""
        print(f"[•] Processing ad-hoc data for organization: {org.name}")
        
        # Validate organization connections
        try:
            org_connections = self.connection_manager.validate_organization_connections(
                org.name, org.api_token_env, org.providers
            )
            print(f"[✓] Validated {len(org_connections)} provider connections")
        except Exception as e:
            return {
                "status": "failed",
                "error": f"Connection validation failed: {str(e)}"
            }
        
        # For ad-hoc uploads, we treat the entire dataset as one "month"
        # This is because ad-hoc data may not have consistent month partitioning
        print(f"[•] Preparing ad-hoc data for upload ({len(df)} rows)")
        
        # Group by provider if cloud_provider column exists
        if "cloud_provider" in df.columns:
            provider_groups = df.group_by("cloud_provider").agg(pl.count().alias("count"))
            print(f"[•] Found providers: {provider_groups.to_dict()}")
        else:
            print("[!] No cloud_provider column found - treating as single provider")
        
        uploaded_files = []
        total_rows = 0
        
        # Process each provider connection
        for provider, connections in org_connections.items():
            if "cloud_provider" in df.columns:
                provider_df = df.filter(pl.col("cloud_provider") == provider)
            else:
                provider_df = df
            
            if provider_df.is_empty():
                print(f"[•] No data for provider {provider}, skipping")
                continue
            
            print(f"[•] Processing {len(provider_df)} rows for provider {provider}")
            
            # For ad-hoc uploads, we use current timestamp as the "month"
            from datetime import datetime
            timestamp = datetime.now().strftime("%Y-%m")
            
            # Format data for S3 upload
            formatted_df = self.s3_formatter.format_for_s3_upload(
                provider_df, 
                provider, 
                timestamp,
                org.name
            )
            
            # Calculate billing period (use current month for ad-hoc)
            billing_period = self._calculate_billing_period(timestamp)
            
            # Upload to S3 for each connection
            for connection in connections:
                try:
                    if not dry_run:
                        result = self.s3_uploader.upload_to_s3(
                            formatted_df,
                            connection.connection_id,
                            billing_period,
                            connection.api_token
                        )
                        
                        uploaded_files.append({
                            "provider": provider,
                            "connection": connection.connection_id,
                            "rows": len(formatted_df),
                            "drop_id": result.get("drop_id"),
                            "s3_key": result.get("s3_key")
                        })
                    else:
                        print(f"[DRY RUN] Would upload {len(formatted_df)} rows to {connection.connection_id}")
                        uploaded_files.append({
                            "provider": provider,
                            "connection": connection.connection_id,
                            "rows": len(formatted_df),
                            "drop_id": "dry-run-drop-id",
                            "s3_key": "dry-run-s3-key"
                        })
                    
                    total_rows += len(formatted_df)
                    
                except Exception as e:
                    print(f"[!] Upload failed for {connection.connection_id}: {str(e)}")
                    return {
                        "status": "failed",
                        "error": f"Upload failed for connection {connection.connection_id}: {str(e)}"
                    }
        
        return {
            "status": "success",
            "files_uploaded": len(uploaded_files),
            "total_rows": total_rows,
            "providers": list(org_connections.keys()),
            "uploaded_files": uploaded_files
        }
    
    def _calculate_billing_period(self, month: str) -> str:
        """
        Calculate billing period from month string.
        
        Args:
            month: Month in YYYY-MM format
            
        Returns:
            Billing period in YYYYMMDD-YYYYMMDD format
        """
        # Parse month
        year, month_num = month.split('-')
        year = int(year)
        month_num = int(month_num)
        
        # Calculate billing period (first day of month to first day of next month)
        start_date = f"{year:04d}{month_num:02d}01"
        
        # Calculate next month
        if month_num == 12:
            next_year = year + 1
            next_month = 1
        else:
            next_year = year
            next_month = month_num + 1
        
        end_date = f"{next_year:04d}{next_month:02d}01"
        
        return f"{start_date}-{end_date}"
    
    def _perform_unified_validation(self, dry_run: bool = False) -> Dict[str, Any]:
        """
        Perform unified validation with stage discovery for both dry-run and actual run.
        
        Validates:
        - S3 bucket access
        - Stage discovery for all provider/organization combinations
        - Connection existence for all organizations
        
        Returns processing plan with discovered stages per provider.
        """
        # Step 1: Validate S3 bucket access
        try:
            self.s3_uploader.validate_bucket_access()
            print("✅ S3 bucket access validated")
        except S3UploadError as e:
            return {
                "status": "failed",
                "error": f"S3 bucket validation failed: {str(e)}",
                "validation_step": "s3_bucket"
            }
        
        # Step 2: Perform stage discovery for all organizations
        try:
            org_discovery_results = self._perform_stage_discovery_for_all_orgs()
            total_files = sum(len(result.available_months) for results in org_discovery_results.values() for result in results)
            total_providers = sum(len(results) for results in org_discovery_results.values())
            
            # Display discovery results
            self._display_stage_discovery_results(org_discovery_results)
            print(f"✅ Stage discovery complete. Processing plan: {total_files} files across {len(self.config.organizations)} organizations")
            
        except (StageDiscoveryError, PartialDataError, NoDataError) as e:
            return {
                "status": "failed",
                "error": f"Stage discovery failed: {str(e)}",
                "validation_step": "stage_discovery"
            }
        except Exception as e:
            return {
                "status": "failed",
                "error": f"Unexpected error during stage discovery: {str(e)}",
                "validation_step": "stage_discovery"
            }
        
        # Step 3: Validate connections for all organizations
        org_validation_results = {}
        for org in self.config.organizations:
            try:
                providers = org.providers  # Now a simple list
                
                if providers:
                    # Get organization API key
                    from ..api.connection_manager import get_org_api_key, ConnectionManager
                    org_api_key = get_org_api_key(org.name)
                    conn_manager = ConnectionManager(org_api_key)
                    
                    print(f"   🔍 Ensuring connections for {len(providers)} providers...")
                    
                    # First check existing connections
                    conn_status = conn_manager.validate_connections_exist(providers)
                    missing_providers = [p for p, exists in conn_status.items() if not exists]
                    
                    if missing_providers:
                        print(f"   ⚙️  Auto-creating {len(missing_providers)} missing connections...")
                        
                        # Prepare S3 config for connection creation
                        s3_config_dict = {
                            "bucket_name": self.s3_config.bucket,
                            "root_folder": self.s3_config.root_folder,
                            "org_folder": org.name  # Use organization name as folder
                        }
                        
                        # Auto-create missing connections
                        try:
                            conn_results = conn_manager.ensure_all_connections(providers, s3_config_dict)
                            
                            # Mark as validated after auto-creation
                            org_validation_results[org.name] = {
                                "status": "validated",
                                "providers": providers,
                                "discovery_results": org_discovery_results[org.name],
                                "auto_created_connections": missing_providers
                            }
                            print(f"   ✅ {org.name}: All connections ready ({len(providers)} providers)")
                        except Exception as auto_create_error:
                            print(f"   ❌ Failed to auto-create connections: {auto_create_error}")
                            org_validation_results[org.name] = {
                                "status": "connection_creation_failed",
                                "missing_providers": missing_providers,
                                "providers": providers,
                                "discovery_results": org_discovery_results[org.name],
                                "error": f"Failed to create connections: {str(auto_create_error)}"
                            }
                    else:
                        print(f"   ✅ {org.name}: All connections validated ({len(providers)} providers)")
                        org_validation_results[org.name] = {
                            "status": "validated",
                            "providers": providers,
                            "discovery_results": org_discovery_results[org.name]
                        }
                else:
                    org_validation_results[org.name] = {
                        "status": "no_providers",
                        "providers": [],
                        "discovery_results": []
                    }
                    
            except Exception as e:
                print(f"   ❌ Exception validating {org.name}: {str(e)}")
                org_validation_results[org.name] = {
                    "status": "failed",
                    "error": str(e),
                    "providers": providers if 'providers' in locals() else [],
                    "discovery_results": org_discovery_results.get(org.name, [])
                }
        
        # Step 4: Generate processing plan
        
        processing_plan = self._generate_processing_plan(org_validation_results, org_discovery_results, dry_run)
        
        
        return {
            "status": "success",
            "dataset_name": self.config.dataset_name,
            "org_validation_results": org_validation_results,
            "processing_plan": processing_plan,
            "discovery_results": org_discovery_results
        }
    
    def _perform_stage_discovery_for_all_orgs(self) -> Dict[str, List]:
        """
        Perform stage discovery for all organizations.
        
        Returns:
            Dict mapping organization name to list of StageDiscoveryResult objects
        """
        org_discovery_results = {}
        
        for org in self.config.organizations:
            # Resolve months for this organization (org-level override or global default)
            org_months = org.months if org.months else self.config.months
            if not org_months:
                raise ValueError(f"Organization {org.name} has no months specified and no global months defined")
            
            # Perform discovery for this organization's providers
            discovery_results = self.stage_discovery.discover_provider_stages(
                dataset_name=self.config.dataset_name,
                providers=org.providers,
                required_months=org_months,
                stage_overrides=org.stage_overrides
            )
            
            org_discovery_results[org.name] = discovery_results
        
        return org_discovery_results
    
    def _display_stage_discovery_results(self, org_discovery_results: Dict[str, List]) -> None:
        """
        Display stage discovery results in a user-friendly format.
        
        Args:
            org_discovery_results: Dict mapping org name to discovery results
        """
        if not org_discovery_results:
            print("No discovery results to display")
            return
        
        print("\n🔍 Stage Discovery Results:")
        
        all_results = []
        for org_name, results in org_discovery_results.items():
            if len(self.config.organizations) > 1:
                print(f"\n   📋 Organization: {org_name}")
                prefix = "     "
            else:
                prefix = "   "
            
            for result in results:
                method_desc = f"({result.discovery_method} - {len(result.available_months)} months available)"
                print(f"{prefix}{result.provider}: {result.stage} {method_desc}")
                all_results.append(result)
        
        # Show overall summary
        discovery_count = sum(1 for r in all_results if r.discovery_method == "discovered")
        override_count = sum(1 for r in all_results if r.discovery_method == "override")
        stages_used = set(r.stage for r in all_results)
        
        print("\nℹ️  Discovery Summary:")
        if override_count > 0:
            print(f"   • {override_count} providers using overrides")
        if discovery_count > 0:
            print(f"   • {discovery_count} providers auto-discovered")
        
        if len(stages_used) > 1:
            print("   • Mixed stages detected - upload will proceed with best available data per provider")
        else:
            stage_name = list(stages_used)[0]
            print(f"   • All providers using {stage_name} stage")
    
    def _generate_processing_plan(self, org_validation_results: Dict[str, Dict], org_discovery_results: Dict[str, List], dry_run: bool) -> Dict[str, Any]:
        """
        Generate processing plan using stage discovery results.
        
        Returns structured plan for streaming processing with discovered stages.
        """
        plan = {
            "organizations": [],
            "total_files": 0,
            "total_organizations": len(self.config.organizations)
        }
        
        for org in self.config.organizations:
            org_result = org_validation_results[org.name]
            discovery_results = org_discovery_results[org.name]
            
            
            if org_result["status"] not in ["validated"]:
                # Skip organizations that failed validation or connection creation
                error_msg = org_result.get('error', f"Status: {org_result['status']}")
                plan["organizations"].append({
                    "name": org.name,
                    "status": org_result["status"],
                    "error": error_msg,
                    "providers": []
                })
                continue
            
            org_plan = {
                "name": org.name,
                "status": "ready",
                "providers": []
            }
            
            
            # Process each provider with its discovered stage
            for discovery_result in discovery_results:
                
                provider_plan = {
                    "name": discovery_result.provider,
                    "stage": discovery_result.stage,
                    "months": discovery_result.available_months,
                    "file_count": len(discovery_result.available_months),
                    "discovery_method": discovery_result.discovery_method
                }
                
                org_plan["providers"].append(provider_plan)
                plan["total_files"] += len(discovery_result.available_months)
            
            plan["organizations"].append(org_plan)
        
        return plan
    
    def _process_file_streaming(
        self, 
        org: Organization, 
        provider: str, 
        month: str, 
        stage: str,
        drop_id: str
    ) -> Dict[str, Any]:
        """
        Process a single file using streaming approach.
        
        Loads one parquet file from discovered stage, processes it through CBF mapping, uploads to S3,
        and releases from memory.
        """
        try:
            # Load single file from S3
            dataset_name = self.config.dataset_name
            s3_key = f"czbillgen/datasets/{dataset_name}/{stage}/{provider}/{month}.parquet"
            s3_path = f"s3://{self.s3_data_service.bucket}/{s3_key}"
            
            # Load DataFrame from S3 using helper method (works around Polars S3 issues)
            df = self.s3_data_service.read_parquet_from_s3(s3_key)
            print(f"      📊 {month}: {len(df):,} rows")
            
            # Apply CBF processing
            processed_data = self.data_processor.process_generated_dataset(df)
            
            # Calculate billing period for this month
            billing_period = self._calculate_billing_period(month)
            
            # Build S3 prefix for this upload
            s3_prefix = self.s3_uploader.build_s3_prefix(
                org_folder=org.name,
                provider=provider,
                billing_period=billing_period,
                drop_id=drop_id
            )
            
            # Get next file number and generate filename
            file_number = self.s3_uploader.get_next_file_number(s3_prefix)
            filename = self.s3_formatter.generate_filename(file_number)
            
            # Upload the processed data to S3 (may result in multiple files)
            upload_results = self.s3_uploader.upload_dataframe_gzip(
                df=processed_data,
                s3_prefix=s3_prefix,
                filename=filename
            )
            
            # Add metadata for manifest updates to each upload result
            for upload_result in upload_results:
                upload_result.update({
                    "provider": provider,
                    "month": month,
                    "billing_period": billing_period,
                    "s3_prefix": s3_prefix,
                    "rows_processed": len(processed_data)
                })
            
            # Calculate total uploaded rows
            total_uploaded_rows = sum(result["rows"] for result in upload_results)
            files_count = len(upload_results)
            
            if files_count == 1:
                print(f"         ✅ Uploaded {total_uploaded_rows:,} rows to S3 (single file)")
            else:
                print(f"         ✅ Uploaded {total_uploaded_rows:,} rows to S3 ({files_count} files)")
            
            # DataFrame will be garbage collected when function returns
            return {
                "status": "success",
                "upload_results": upload_results,  # Changed from single to list
                "files_count": files_count,
                "total_rows": total_uploaded_rows
            }
            
        except Exception as e:
            print(f"         ❌ Failed to process {provider}/{month}: {str(e)}")
            return {
                "status": "failed",
                "error": str(e),
                "provider": provider,
                "month": month
            }
    
    def _stream_process_organization(
        self, 
        org: Organization, 
        org_plan: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Process an organization using streaming approach.
        
        Processes files one at a time to minimize memory usage.
        """
        print(f"   📤 STREAMING: Processing data for {org.name}")
        
        # Generate drop ID for this upload
        drop_id = self.s3_formatter.generate_drop_id()
        print(f"   🆔 Drop ID: {drop_id}")
        
        uploaded_files = []
        failed_files = []
        total_rows = 0
        
        # Process each provider
        for provider_plan in org_plan["providers"]:
            provider = provider_plan["name"]
            stage = provider_plan["stage"]
            months = provider_plan["months"]
            
            print(f"   🌤️  Processing provider: {provider} (from {stage} stage)")
            print(f"      📅 Processing {len(months)} months: {months}")
            
            # Process each month file individually
            for month in months:
                file_result = self._process_file_streaming(org, provider, month, stage, drop_id)
                
                if file_result["status"] == "success":
                    # Handle multiple upload results per file
                    upload_results = file_result["upload_results"]
                    uploaded_files.extend(upload_results)
                    total_rows += file_result["total_rows"]
                else:
                    failed_files.append(file_result)
        
        # Update manifests and prune old drops
        if uploaded_files:
            try:
                self._update_manifests_and_prune(org, uploaded_files, drop_id)
            except Exception as e:
                print(f"      ⚠️  Manifest update failed: {str(e)}")
                # Don't fail the upload for manifest issues
        
        # Determine result status
        if failed_files and not uploaded_files:
            status = "failed"
        elif failed_files and uploaded_files:
            status = "partial"
        else:
            status = "success"
        
        return {
            "status": status,
            "organization": org.name,
            "drop_id": drop_id,
            "files_uploaded": len(uploaded_files),
            "files_failed": len(failed_files),
            "total_rows": total_rows,
            "providers": [p["name"] for p in org_plan["providers"]],
            "uploaded_files": uploaded_files,
            "failed_files": failed_files
        }
    
    
    def _update_manifests_and_prune(
        self, 
        org: Organization, 
        uploaded_files: List[Dict[str, Any]], 
        drop_id: str
    ) -> None:
        """Update manifest files and prune old drops."""
        # Group uploaded files by provider and billing period for manifest updates
        manifest_updates = {}
        
        for file_info in uploaded_files:
            provider = file_info["provider"]
            billing_period = file_info["billing_period"]
            
            key = (provider, billing_period)
            if key not in manifest_updates:
                manifest_updates[key] = []
            manifest_updates[key].append(file_info)
        
        # Update manifest for each provider/period combination
        for (provider, billing_period), files in manifest_updates.items():
            try:
                # Update manifest
                manifest_prefix = self.s3_uploader.build_manifest_prefix(
                    org_folder=org.name,
                    provider=provider,
                    billing_period=billing_period
                )
                manifest_data = self.s3_formatter.create_manifest_data(drop_id)
                self.s3_uploader.upload_manifest(manifest_data, manifest_prefix)
                
                # Prune old drops (keep latest 3)
                self.s3_uploader.prune_old_drops(
                    org_folder=org.name,
                    provider=provider,
                    billing_period=billing_period,
                    max_drops=3
                )
                
            except S3UploadError as e:
                print(f"      ⚠️  Manifest/pruning warning for {provider}/{billing_period}: {str(e)}")
                # Don't fail the entire upload for manifest issues
    
    def show_recreate_plan(self) -> Dict[str, Any]:
        """
        Show what would be deleted/recreated in recreate-connections dry-run mode.
        
        Returns:
            Plan summary for display
        """
        print(f"🔍 Recreate Connections Analysis for dataset: {self.config.dataset_name}")
        print("=" * 70)
        
        # Generate recreation plan
        plan = self.connection_recreator.generate_recreation_plan(self.config.organizations)
        
        # Display connections to delete
        if plan["total_connections"] > 0:
            print(f"\n❌ CloudZero Connections to Delete:")
            for org_plan in plan["organizations"]:
                for conn in org_plan["connections"]:
                    print(f"   ✗ {conn['id']} ({conn['provider']})")
        else:
            print(f"\n✅ No CloudZero connections to delete")
        
        # Display S3 folders to delete
        if plan["total_folders"] > 0:
            print(f"\n🗂️  S3 Org Folders to Delete:")
            for org_plan in plan["organizations"]:
                for folder in org_plan["folders"]:
                    print(f"   ✗ {folder['prefix']} ({folder['objects']} files, {folder['size_formatted']})")
        else:
            print(f"\n✅ No S3 folders to delete")
        
        # Display preserved data
        print(f"\n💾 Source Data Preserved:")
        print(f"   ✅ s3://{self.s3_config.bucket}/czbillgen/{self.config.dataset_name}/ (all datasets intact)")
        
        # Display connections to recreate
        if plan["total_connections"] > 0:
            print(f"\n🔄 Connections to Recreate:")
            for org_plan in plan["organizations"]:
                org_name = org_plan["name"]
                providers = set(conn['provider'] for conn in org_plan["connections"])
                for provider in providers:
                    print(f"   ✅ {org_name}/{provider} → new first ingest performance")
        
        # Display impact summary
        print(f"\n📊 Recreation Impact:")
        print(f"   Total files to delete: {plan['total_objects']:,}")
        print(f"   Total size to delete: {plan['total_size_formatted']}")
        print(f"   Estimated duration: ~{plan['estimated_duration_minutes']:.0f} minutes")
        print(f"   Organizations affected: {len(plan['organizations'])}")
        print(f"   Connections to recreate: {plan['total_connections']}")
        
        return {
            "status": "success",
            "mode": "recreate_connections_dry_run",
            "dataset": self.config.dataset_name,
            "plan": plan,
            "organizations_processed": len(plan["organizations"]),
            "organizations_successful": len(plan["organizations"]),
            "results": {"dry_run_plan": plan}
        }
    
    def recreate_and_upload(self) -> Dict[str, Any]:
        """
        Execute connection recreation followed by normal upload.
        
        Returns:
            Combined results from recreation and upload
        """
        print(f"🔄 Starting connection recreation and upload for dataset: {self.config.dataset_name}")
        
        recreation_results = []
        successful_recreations = 0
        
        # Phase 1: Recreate connections for each organization
        print(f"\n🔄 Phase 1: Recreating connections...")
        for i, org in enumerate(self.config.organizations, 1):
            print(f"\n  Organization {i}/{len(self.config.organizations)}: {org.name}")
            
            result = self.connection_recreator.recreate_org_connections(org, dry_run=False)
            recreation_results.append(result)
            
            if result.status == "success":
                successful_recreations += 1
                print(f"    ✅ Recreation successful")
                print(f"       • Deleted {len(result.deleted_connections)} connections")
                print(f"       • Deleted {sum(f.objects_deleted for f in result.deleted_folders)} S3 objects")
                print(f"       • Created {len(result.created_connections)} connections")
            elif result.status == "partial":
                print(f"    ⚠️  Recreation partially successful: {result.error_message}")
            else:
                print(f"    ❌ Recreation failed: {result.error_message}")
        
        # Display recreation summary
        recreation_summary = self.connection_recreator.format_recreation_summary(recreation_results)
        print(f"\n📊 Recreation Summary:")
        print(f"   Status: {recreation_summary['overall_status']}")
        print(f"   Organizations: {recreation_summary['successful_orgs']}/{recreation_summary['total_orgs']} successful")
        print(f"   Connections recreated: {recreation_summary['total_connections_created']}")
        print(f"   Objects deleted: {recreation_summary['total_objects_deleted']:,}")
        print(f"   Size deleted: {recreation_summary['total_size_formatted']}")
        
        # Phase 2: Execute normal upload workflow
        if recreation_summary['overall_status'] != "failed":
            print(f"\n🚀 Phase 2: Uploading data with recreated connections...")
            
            # Initialize data processor for upload
            self.data_processor = DataProcessor()
            
            upload_result = self.upload_dataset_to_orgs(dry_run=False)
            
            # Combine results
            return {
                "status": upload_result["status"],
                "mode": "recreate_and_upload",
                "dataset": self.config.dataset_name,
                "organizations_processed": len(self.config.organizations),
                "organizations_successful": successful_recreations,
                "recreation_summary": recreation_summary,
                "upload_result": upload_result,
                "results": {
                    "recreation": recreation_results,
                    "upload": upload_result.get("results", {})
                }
            }
        else:
            # Recreation failed, don't attempt upload
            return {
                "status": "failed",
                "mode": "recreate_and_upload",
                "dataset": self.config.dataset_name,
                "error": "Connection recreation failed, upload aborted",
                "organizations_processed": len(self.config.organizations),
                "organizations_successful": 0,
                "recreation_summary": recreation_summary,
                "results": {
                    "recreation": recreation_results
                }
            }
    
