"""
Streaming generation engine for memory-efficient billing data generation.

Provides job-level streaming with memory bounds, reference data filtering optimization,
and line item tracking for CloudZero ingestion cost planning.
"""

from datetime import datetime, timedelta
from pathlib import Path
from typing import List, Dict, Optional, Any, Tuple
from dataclasses import asdict
import logging
import uuid
import polars as pl
from rich.console import Console
from rich.progress import SpinnerColumn, TextColumn, BarColumn, TimeElapsedColumn

from ..config.models import GenerateConfig, GenerationJob
from ..core import (
    apply_filters,
    load_reference_data_for_job,
    estimate_record_count_after_filtering,
    generate_resources,
    generate_resources_batch,
    merge_resources_with_time_series,
    JobSplitter,
    LineItemTracker,
    calculate_hours_in_date_range,
)
from ..processing import (
    apply_variance,
    apply_daily_spike,
    apply_daily_step_up,
    apply_daily_step_down,
    apply_hourly_spike,
    apply_hourly_step_up,
    apply_hourly_step_down,
    DEFAULT_VARIANCE,
    DEFAULT_ANOMALY,
)
from ..processing.step_state_manager import initialize_step_state_manager, cleanup_step_state
from ..algorithms.jitter import apply_resource_level_jitter
from ..validation import (
    validate_basic_completeness,
    validate_cost_variance, 
    validate_cost_conservation,
    ValidationError
)
from src.services.s3_data_service import S3DataService
from src.services.optimized_s3_service import OptimizedS3Service, StreamingUploadManager
from src.services.progress_wrapper import create_progress
from .bounded_cache import BoundedReferenceCache


class GenerationError(Exception):
    """Raised when generation fails."""
    pass


class StreamingGenerationEngine:
    """
    Memory-efficient streaming generation engine.
    
    Features:
    - Job-level streaming with memory bounds
    - Reference data filtering optimization
    - Line item tracking for cost planning
    - Progress reporting and error handling
    """
    
    def __init__(
        self,
        s3_service: S3DataService,
        console: Console,
        target_memory_bytes: int = 1024 * 1024 * 1024,  # 1GB default
        bytes_per_record: int = 151,
        org_name: str = "default",
        use_optimized_s3: bool = True,
        use_vectorized_operations: bool = True
    ):
        # Use optimized S3 service if available and requested
        if use_optimized_s3 and not isinstance(s3_service, OptimizedS3Service):
            self.s3_service = OptimizedS3Service()
            self.upload_manager = StreamingUploadManager(self.s3_service)
        else:
            self.s3_service = s3_service
            self.upload_manager = None

        self.console = console
        self.logger = logging.getLogger(__name__)
        self.target_memory_bytes = target_memory_bytes
        self.bytes_per_record = bytes_per_record
        self.org_name = org_name
        self.use_optimized_s3 = use_optimized_s3
        self.use_vectorized_operations = use_vectorized_operations
        
        # Initialize components
        self.job_splitter = JobSplitter(target_memory_bytes, bytes_per_record)
        self.line_item_tracker = LineItemTracker()
        
        # Initialize bounded reference data cache (25% of memory budget)
        cache_memory_budget = target_memory_bytes // 4
        self.reference_data_cache = BoundedReferenceCache(max_memory_bytes=cache_memory_budget)
        
        # Data accumulation for preventing overwrite issues
        self.accumulated_data = {}  # Dict[(provider, month)] -> List[pl.DataFrame]
        self.current_accumulation_size = 0
        self._load_generation_config()
        
        # Generation state
        self.current_dataset_name: Optional[str] = None
        self.failed_jobs: List[str] = []
    
    def _load_generation_config(self):
        """Load generation configuration from environment variables."""
        from src.services.environment_service import env_service
        config = env_service.get_generation_config()
        
        # Convert MB to bytes for internal use
        self.accumulation_memory_limit = config["accumulation_memory_mb"] * 1024 * 1024
        self.memory_log_level = config["memory_log_level"]
        
        # Set up log level mapping
        self._log_levels = {"DEBUG": 0, "INFO": 1, "WARN": 2, "ERROR": 3}
        self._current_log_level = self._log_levels.get(self.memory_log_level, 1)
    
    def _should_log_memory(self, level: str) -> bool:
        """Check if memory message should be logged based on configuration."""
        message_level = self._log_levels.get(level.upper(), 1)
        return message_level >= self._current_log_level
    
    def _estimate_dataframe_memory(self, df: pl.DataFrame) -> int:
        """Estimate memory usage of a DataFrame in bytes."""
        return df.estimated_size()
    
    def generate_dataset(
        self,
        config: GenerateConfig,
        dataset_name: str,
        verbose: bool = False
    ) -> Dict[str, Any]:
        """
        Generate a complete dataset from configuration.
        
        Args:
            config: Generation configuration
            dataset_name: Name for the dataset
            verbose: Enable verbose logging
            
        Returns:
            Dictionary with generation results and statistics
        """
        # Initialize step state manager for cross-month anomaly persistence
        job_uuid = str(uuid.uuid4())
        initialize_step_state_manager(job_uuid)
        
        try:
            self.current_dataset_name = dataset_name
            self.line_item_tracker.reset()
            self.failed_jobs.clear()
            
            # Phase 1: Analyze jobs and estimate memory usage
            self.console.print("[blue]📊 Phase 1: Analyzing generation jobs...[/blue]")
            analysis_results = self._analyze_jobs(config.generation_jobs, verbose)
            
            # Phase 2: Split jobs if needed for memory bounds
            self.console.print("[blue]🔀 Phase 2: Optimizing jobs for memory constraints...[/blue]")
            final_jobs = self._split_jobs_for_memory(
                config.generation_jobs, 
                analysis_results["reference_matches_per_job"],
                verbose
            )
            
            # Phase 3: Generate data for each job
            self.console.print("[blue]⚡ Phase 3: Generating billing data...[/blue]")
            generation_results = self._generate_jobs(final_jobs, verbose)
            
            # Phase 4: Collect statistics and results
            self.console.print("[blue]📈 Phase 4: Collecting generation results...[/blue]")
            results = self._collect_results(generation_results, analysis_results)
            
            # Phase 5: Finalize generation by uploading any remaining accumulated data
            self.console.print("[blue]🏁 Phase 5: Finalizing generation (consolidating temp files)...[/blue]")
            self.finalize_generation()
            
            # Phase 6: Write success marker to indicate complete generation
            self.console.print("[blue]✅ Phase 6: Writing success marker...[/blue]")
            self._write_success_marker(dataset_name)
            
            return results
            
        finally:
            # Always clean up step state manager, even on failures
            self.console.print("[blue]🧹 Cleaning up cross-month step state...[/blue]")
            cleanup_step_state()
    
    def _analyze_jobs(
        self,
        jobs: List[GenerationJob],
        verbose: bool = False
    ) -> Dict[str, Any]:
        """Analyze jobs and estimate memory requirements."""
        self.console.print("[blue]🔍 Analyzing generation jobs...[/blue]")
        
        reference_matches_per_job = {}
        total_estimated_records = 0
        
        with create_progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            console=self.console,
            transient=not verbose
        ) as progress:
            task = progress.add_task("Analyzing jobs...", total=len(jobs))
            
            for job in jobs:
                if verbose:
                    progress.update(task, description=f"Analyzing job: {job.name}")
                
                # Load reference data optimized for this job
                providers = list(set(f.cloud_provider for f in job.filters))
                services = []
                for filter_obj in job.filters:
                    services.extend(s.name for s in filter_obj.services)
                
                ref_df = self._load_reference_data_for_job(
                    providers[0], services
                )
                
                # Apply filters and estimate record count
                filtered_results = apply_filters(ref_df, [asdict(f) for f in job.filters])
                reference_matches = sum(len(filtered_df) for filtered_df, _ in filtered_results)
                
                hours_in_range = calculate_hours_in_date_range(job.start_date, job.end_date)
                estimated_records = reference_matches * job.resource_count * hours_in_range
                
                reference_matches_per_job[job.name] = reference_matches
                total_estimated_records += estimated_records
                
                if verbose:
                    self.console.print(f"  [cyan]{job.name}[/cyan]: {reference_matches} ref matches, {estimated_records:,} estimated records")
                
                progress.update(task, advance=1)
        
        total_estimated_memory = total_estimated_records * self.bytes_per_record
        
        self.console.print(f"[green]✓ Analysis complete: {total_estimated_records:,} total records, {total_estimated_memory / (1024**3):.2f} GB estimated[/green]")
        
        return {
            "reference_matches_per_job": reference_matches_per_job,
            "total_estimated_records": total_estimated_records,
            "total_estimated_memory": total_estimated_memory,
            "jobs_analyzed": len(jobs)
        }
    
    def _split_jobs_for_memory(
        self,
        jobs: List[GenerationJob],
        reference_matches_per_job: Dict[str, int],
        verbose: bool = False
    ) -> List[GenerationJob]:
        """Split jobs that exceed memory bounds."""
        self.console.print("[blue]🔄 Checking memory bounds and splitting jobs...[/blue]")
        
        split_jobs = self.job_splitter.split_jobs_for_memory_bounds(
            jobs, reference_matches_per_job
        )
        
        if len(split_jobs) > len(jobs):
            splits_created = len(split_jobs) - len(jobs)
            self.console.print(f"[yellow]⚠️  Created {splits_created} job splits for memory bounds[/yellow]")
            
            if verbose:
                for i, job in enumerate(split_jobs):
                    if "_split_" in job.name:
                        self.console.print(f"  [cyan]{job.name}[/cyan]: {job.start_date} to {job.end_date}")
        else:
            self.console.print("[green]✓ All jobs fit within memory bounds[/green]")
        
        return split_jobs
    
    def _generate_jobs(
        self,
        jobs: List[GenerationJob],
        verbose: bool = False
    ) -> Dict[str, Any]:
        """Generate data for all jobs."""
        self.console.print(f"[blue]🚀 Generating data for {len(jobs)} job(s)...[/blue]")
        
        successful_jobs = 0
        failed_jobs = 0
        total_records_generated = 0
        total_actual_cost = 0.0
        job_costs = {}

        with create_progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            TimeElapsedColumn(),
            console=self.console,
            transient=not verbose
        ) as progress:
            task = progress.add_task("Generating jobs...", total=len(jobs))
            
            for job in jobs:
                if verbose:
                    progress.update(task, description=f"Generating: {job.name}")
                
                # Generate data for this job - any failure will immediately stop everything
                job_results = self._generate_single_job(job, verbose)
                
                # Track results
                total_records_generated += job_results["records_generated"]
                successful_jobs += 1

                # Track actual cost if available
                if "actual_cost" in job_results:
                    job_actual_cost = job_results["actual_cost"]
                    total_actual_cost += job_actual_cost
                    job_costs[job.name] = job_actual_cost
                
                if verbose:
                    self.console.print(f"  [green]✓ {job.name}[/green]: {job_results['records_generated']:,} records")
                
                progress.update(task, advance=1)
        
        self.console.print(f"[green]✓ Generation complete: {total_records_generated:,} records in {successful_jobs} job(s)[/green]")
        
        return {
            "successful_jobs": successful_jobs,
            "failed_jobs": failed_jobs,
            "total_records_generated": total_records_generated,
            "failed_job_names": self.failed_jobs,
            "total_actual_cost": total_actual_cost,
            "job_costs": job_costs
        }
    
    def _generate_single_job(
        self,
        job: GenerationJob,
        verbose: bool = False
    ) -> Dict[str, Any]:
        """Generate data for a single job with fail-fast error handling."""
        # Load reference data optimized for this job
        providers = list(set(f.cloud_provider for f in job.filters))
        services = []
        for filter_obj in job.filters:
            services.extend(s.name for s in filter_obj.services)
        
        ref_df = self._load_reference_data_for_job(
            providers[0], services
        )

        # Sort reference data for deterministic filtering and resource generation
        # This ensures consistent resource IDs regardless of filter batch boundaries
        ref_df = ref_df.with_columns([
            pl.col('usage_family').fill_null(""),
            pl.col('operation').fill_null(""),
            pl.col('usage_type').fill_null("")
        ]).sort([
            'cloud_provider', 'service', 'resource_type',
            'usage_family', 'operation', 'usage_type'
        ])

        # Apply filters to reference data
        filtered_results = apply_filters(ref_df, [asdict(f) for f in job.filters])
        
        if not filtered_results:
            return {"records_generated": 0, "months_generated": []}
        
        # Calculate total number of resources that will be generated
        total_resources = 0
        for filtered_df, _ in filtered_results:
            total_resources += len(filtered_df) * job.resource_count
        
        # Generate time series data with correct resource count
        time_series_df = self._generate_time_series(job, total_resources)
        
        # Generate resources and merge with time series with streaming validation
        total_records = 0
        months_generated = []
        job_cost_accumulator = 0.0  # Track total cost for conservation validation
        
        for filtered_df, filter_config in filtered_results:
            # Generate resources for this filtered subset
            resources_df = self._generate_resources_for_subset(
                filtered_df, job.resource_count, ref_df, job.uuid
            )
            
            # Apply variance and anomaly profiles
            enhanced_time_series = self._apply_profiles(
                time_series_df, filter_config, job
            )
            
            # Merge resources with time series
            merged_df = merge_resources_with_time_series(resources_df, enhanced_time_series)
            
            # Apply resource-level jitter
            # For now, use a simple default jitter config
            # In a full implementation, this would be configurable
            jitter_config = {
                "base_deviation": 0.15,
                "time_correlation": 0.7, 
                "resource_uniqueness": 0.8
            }
            
            # Apply jitter using vectorized operations by resource
            final_df = self._apply_jitter_to_costs(merged_df, jitter_config)
            
            # Streaming validation - validate chunk immediately without accumulating
            self._validate_data_chunk(final_df, job)
            
            # Accumulate cost for final conservation check
            if "cost" in final_df.columns:
                job_cost_accumulator += float(final_df["cost"].sum())
            
            # Upload to S3 by month
            months_in_job = self._upload_by_month(final_df, providers[0], job)
            months_generated.extend(months_in_job)
            
            total_records += len(final_df)
            
            # Force garbage collection if memory usage is high
            self._check_memory_pressure_and_gc()
        
        # Final cost conservation validation (without accumulating data)
        if job_cost_accumulator > 0 and abs(job_cost_accumulator - job.total_cost) > 0.01:
            cost_variance = abs(job_cost_accumulator - job.total_cost) / job.total_cost * 100
            if cost_variance > 5.0:  # Allow 5% variance
                self.console.print(f"⚠️ Cost conservation warning: Generated ${job_cost_accumulator:.2f}, Expected ${job.total_cost:.2f} ({cost_variance:.1f}% variance)")
            else:
                self.console.print(f"✓ Cost conservation OK: Generated ${job_cost_accumulator:.2f}, Expected ${job.total_cost:.2f} ({cost_variance:.1f}% variance)")
        
        return {
            "records_generated": total_records,
            "months_generated": months_generated,
            "actual_cost": job_cost_accumulator
        }
    
    def _load_reference_data_for_job(
        self,
        provider: str,
        services: List[str]
    ) -> pl.DataFrame:
        """Load reference data with caching using S3DataService."""
        cache_key = f"{provider}:{':'.join(sorted(services))}"
        
        # Check bounded cache first
        cached_data = self.reference_data_cache.get(cache_key)
        if cached_data is not None:
            return cached_data
        
        # Use S3DataService to get reference data path automatically
        ref_data_path = self.s3_service.ensure_reference_data()
        
        # Use the existing core function to load and filter reference data
        reference_data = load_reference_data_for_job(ref_data_path, provider, services)
        
        # Store in bounded cache
        self.reference_data_cache.put(cache_key, reference_data)
        
        # Proactively check for memory pressure
        self.reference_data_cache.evict_if_memory_pressure()
        
        return reference_data
    
    def _generate_time_series(self, job: GenerationJob, total_resources: int) -> pl.DataFrame:
        """Generate time series data for a job."""
        # Generate hourly timestamps
        hours = calculate_hours_in_date_range(job.start_date, job.end_date)
        timestamps = []
        current_time = job.start_date
        
        for _ in range(hours):
            timestamps.append(current_time)
            current_time += timedelta(hours=1)
        
        # Calculate base hourly cost per resource
        # Total cost must be distributed across all resources and all hours
        base_hourly_cost = job.total_cost / (hours * total_resources)
        
        return pl.DataFrame({
            "timestamp": timestamps,
            "hourly_cost": [base_hourly_cost] * len(timestamps)
        })
    
    def _generate_resources_for_subset(
        self,
        filtered_df: pl.DataFrame,
        resource_count: int,
        ref_df: pl.DataFrame,
        job_uuid: str = None
    ) -> pl.DataFrame:
        """Generate resources for a filtered subset of reference data."""
        all_resources = []
        
        # Use vectorized batch processing by default for better performance
        if not filtered_df.is_empty():
            resources = generate_resources_batch(filtered_df, resource_count, ref_df, self.org_name, job_uuid)
            all_resources.append(resources)
        
        if all_resources:
            return pl.concat(all_resources, how="vertical")
        else:
            return pl.DataFrame()
    
    def _apply_profiles(
        self,
        time_series_df: pl.DataFrame,
        filter_config: Dict[str, Any],
        job: GenerationJob
    ) -> pl.DataFrame:
        """Apply variance and anomaly profiles to time series data."""
        enhanced_df = time_series_df.clone()
        
        # Generate resource_id for deterministic processing
        provider = filter_config.get("cloud_provider", "unknown")
        resource_id = f"{job.name}_{provider}"
        
        # Apply variance profile (use defaults if not specified or empty)
        variance_config = filter_config.get("variance_profile", DEFAULT_VARIANCE)
        if not variance_config:  # Handle empty string case
            variance_config = DEFAULT_VARIANCE
        enhanced_df = apply_variance(enhanced_df, variance_config, resource_id)
        
        # Apply anomaly profile (optional - no anomaly by default)
        anomaly_config = filter_config.get("anomaly_profile")
        if anomaly_config:
            # Add variance_config to anomaly_config for auto-calculation if not explicitly provided
            enhanced_anomaly_config = anomaly_config.copy()
            if "variance_config" not in enhanced_anomaly_config and "intensity" not in enhanced_anomaly_config:
                enhanced_anomaly_config["variance_config"] = variance_config
            
            anomaly_type = enhanced_anomaly_config.get("type")
            
            # Route to appropriate anomaly function based on type
            if anomaly_type == "daily_spike":
                enhanced_df = apply_daily_spike(enhanced_df, enhanced_anomaly_config, resource_id)
            elif anomaly_type == "daily_step_up":
                enhanced_df = apply_daily_step_up(enhanced_df, enhanced_anomaly_config, resource_id)
            elif anomaly_type == "daily_step_down":
                enhanced_df = apply_daily_step_down(enhanced_df, enhanced_anomaly_config, resource_id)
            elif anomaly_type == "hourly_spike":
                enhanced_df = apply_hourly_spike(enhanced_df, enhanced_anomaly_config, resource_id)
            elif anomaly_type == "hourly_step_up":
                enhanced_df = apply_hourly_step_up(enhanced_df, enhanced_anomaly_config, resource_id)
            elif anomaly_type == "hourly_step_down":
                enhanced_df = apply_hourly_step_down(enhanced_df, enhanced_anomaly_config, resource_id)
            else:
                self.console.print(f"[yellow]Warning: Unknown anomaly type '{anomaly_type}', skipping[/yellow]")
        
        return enhanced_df

    def _apply_jitter_to_costs(
        self,
        df: pl.DataFrame,
        jitter_config: Dict[str, Any]
    ) -> pl.DataFrame:
        """
        Apply resource-level jitter to cost data using group-by operations.

        Processes each resource independently to apply jitter based on
        resource-specific characteristics and temporal patterns.

        Args:
            df: DataFrame with cost and timestamp columns
            jitter_config: Jitter configuration dictionary

        Returns:
            DataFrame with jittered costs
        """
        if df.is_empty():
            return df

        # Use partition_by for efficient grouped processing
        # This avoids explicit iteration while maintaining resource isolation
        partitioned = df.partition_by(["resource_id"], as_dict=True, maintain_order=True)

        jittered_parts = []
        for (resource_id,), resource_df in partitioned.items():
            # Apply jitter algorithm per resource
            # Note: jitter algorithm requires Python lists as it uses complex
            # temporal correlation logic that's difficult to vectorize in Polars
            jittered_df = self._apply_jitter_to_resource(
                resource_df, resource_id, jitter_config
            )
            jittered_parts.append(jittered_df)

        # Efficient concatenation of results
        return pl.concat(jittered_parts, how="vertical") if jittered_parts else df

    def _apply_jitter_to_resource(
        self,
        resource_df: pl.DataFrame,
        resource_id: str,
        jitter_config: Dict[str, Any]
    ) -> pl.DataFrame:
        """
        Apply jitter to a single resource's cost data.

        This method interfaces with the external jitter algorithm which
        requires Python lists for its temporal correlation calculations.

        Args:
            resource_df: DataFrame for single resource
            resource_id: Resource identifier
            jitter_config: Jitter configuration

        Returns:
            DataFrame with jittered costs for this resource
        """
        # Extract to Python for jitter algorithm (required by external API)
        costs = resource_df["cost"].to_numpy()
        timestamps = resource_df["timestamp"].cast(pl.Utf8).to_numpy()

        # Apply external jitter algorithm
        jittered_costs = apply_resource_level_jitter(
            costs.tolist(),
            timestamps.tolist(),
            resource_id,
            jitter_config
        )

        # Update costs in DataFrame
        return resource_df.with_columns([
            pl.Series("cost", jittered_costs)
        ])

    def _upload_by_month(
        self,
        df: pl.DataFrame,
        provider: str,
        job: GenerationJob
    ) -> List[str]:
        """Accumulate data by month instead of immediate upload to prevent overwriting."""
        if df.is_empty():
            return []
        
        # Group by month
        df_with_month = df.with_columns(
            pl.col("timestamp").dt.strftime("%Y-%m").alias("month")
        )
        
        months_processed = []
        
        # Accumulate data by (provider, month) combination
        for month in df_with_month["month"].unique().to_list():
            month_df = df_with_month.filter(pl.col("month") == month).drop("month")
            
            if month_df.is_empty():
                continue
            
            # Estimate memory usage of this chunk
            chunk_memory = self._estimate_dataframe_memory(month_df)
            
            # Check if adding this chunk would exceed memory limit
            if self.current_accumulation_size + chunk_memory > self.accumulation_memory_limit:
                if self._should_log_memory("INFO"):
                    memory_mb = self.current_accumulation_size / (1024 * 1024)
                    limit_mb = self.accumulation_memory_limit / (1024 * 1024)
                    self.console.print(f"[yellow]💾 Memory threshold reached ({memory_mb:.1f}MB/{limit_mb:.1f}MB), flushing accumulated data[/yellow]")
                self._flush_accumulated_data()
            
            # Add to accumulation
            key = (provider, month)
            if key not in self.accumulated_data:
                self.accumulated_data[key] = []
            
            self.accumulated_data[key].append(month_df)
            self.current_accumulation_size += chunk_memory
            
            # Track line items for reporting (tracking happens during generation)
            self.line_item_tracker.track_dataframe(month_df, provider, month)
            months_processed.append(month)
            
            if self._should_log_memory("DEBUG"):
                rows = len(month_df)
                memory_mb = chunk_memory / (1024 * 1024)
                total_mb = self.current_accumulation_size / (1024 * 1024)
                self.console.print(f"[dim]📊 Accumulated {rows:,} rows for {provider}/{month} ({memory_mb:.1f}MB, total: {total_mb:.1f}MB)[/dim]")
        
        return months_processed
    
    def _flush_accumulated_data(self):
        """Merge accumulated data and upload to temporary S3 files."""
        if not self.accumulated_data:
            return
        
        temp_uploads = []
        total_rows = 0
        
        # Create timestamp for this flush batch
        from datetime import datetime
        flush_timestamp = datetime.now().strftime("%H%M%S")
        
        # Merge all DataFrames for each (provider, month) combination
        for (provider, month), dataframes in self.accumulated_data.items():
            if dataframes:
                # Merge all DataFrames for this provider/month
                if len(dataframes) == 1:
                    merged_df = dataframes[0]
                else:
                    merged_df = pl.concat(dataframes, how="vertical")
                
                # Create temp file entry with timestamp
                temp_uploads.append((self.current_dataset_name, provider, month, merged_df, flush_timestamp))
                total_rows += len(merged_df)
                
                if self._should_log_memory("DEBUG"):
                    self.console.print(f"[dim]📤 Prepared {len(merged_df):,} rows for temp upload: {provider}/{month}_temp_{flush_timestamp}[/dim]")
        
        # Perform batch upload to temp files
        if temp_uploads:
            if self._should_log_memory("INFO"):
                memory_mb = self.current_accumulation_size / (1024 * 1024)
                self.console.print(f"[blue]📤 Uploading {len(temp_uploads)} temp files ({total_rows:,} total rows, {memory_mb:.1f}MB)[/blue]")
            
            self._perform_temp_batch_upload(temp_uploads)
        
        # Clear accumulation and reset memory counter only after successful upload
        self.accumulated_data.clear()
        self.current_accumulation_size = 0
    
    def _perform_temp_batch_upload(self, temp_uploads):
        """Perform the actual batch upload to temporary S3 files with fail-fast error handling."""
        # Convert temp uploads to format expected by S3 service
        converted_uploads = []
        for dataset_name, provider, month, month_df, flush_timestamp in temp_uploads:
            # Use temp file naming pattern: {month}_temp_{timestamp}
            temp_month = f"{month}_temp_{flush_timestamp}"
            converted_uploads.append((dataset_name, provider, temp_month, month_df))
        
        # Use existing batch upload infrastructure
        if self.use_optimized_s3 and isinstance(self.s3_service, OptimizedS3Service):
            upload_result = self.s3_service.upload_multiple_datasets_batch(converted_uploads)
            if not upload_result["success"]:
                raise GenerationError(f"Temp batch upload failed: {upload_result['errors']}")
        else:
            # Fall back to sequential uploads
            for dataset_name, provider, temp_month, month_df in converted_uploads:
                self.s3_service.upload_dataset(dataset_name, provider, temp_month, month_df)
    
    def _perform_batch_upload(self, month_uploads):
        """Perform the actual batch upload to S3 with fail-fast error handling."""
        # Use batch upload if optimized S3 service is available
        if self.use_optimized_s3 and isinstance(self.s3_service, OptimizedS3Service):
            upload_result = self.s3_service.upload_multiple_datasets_batch(month_uploads)
            if not upload_result["success"]:
                raise GenerationError(f"Batch upload failed: {upload_result['errors']}")
        else:
            # Fall back to sequential uploads
            for dataset_name, provider, month, month_df in month_uploads:
                self.s3_service.upload_dataset(dataset_name, provider, month, month_df)
    
    def finalize_generation(self):
        """Finalize generation: flush remaining data and consolidate temp files."""
        
        # Step 1: Upload any remaining accumulated data as temp files
        if self.accumulated_data:
            if self._should_log_memory("INFO"):
                total_keys = len(self.accumulated_data)
                memory_mb = self.current_accumulation_size / (1024 * 1024)
                self.console.print(f"[blue]🏁 Flushing remaining accumulated data: {total_keys} datasets ({memory_mb:.1f}MB)[/blue]")
            
            self._flush_accumulated_data()
        
        
        # Step 2: Discover ALL temp files (including those just created by flush)
        temp_file_groups = self._discover_temp_files()
        
        
        # Step 3: Consolidate using the fresh snapshot
        if self._should_log_memory("INFO"):
            self.console.print(f"[blue]🔄 Consolidating temporary files into final datasets[/blue]")
        
        self._consolidate_temp_files(temp_file_groups)
    
    def _consolidate_temp_files(self, temp_file_groups: Dict[tuple, List[str]]):
        """Consolidate all temporary files into final dataset files."""
        
        if not temp_file_groups:
            if self._should_log_memory("DEBUG"):
                self.console.print("[dim]No temporary files found to consolidate[/dim]")
            return
        
        total_consolidations = len(temp_file_groups)
        
        if self._should_log_memory("INFO"):
            self.console.print(f"[blue]📋 Found {total_consolidations} provider/month combinations to consolidate[/blue]")
        
        loop_count = 0
        try:
            for (provider, month), temp_files in temp_file_groups.items():
                loop_count += 1
                self._consolidate_provider_month(provider, month, temp_files)
                
        except Exception as e:
            self.logger.error(f"Consolidation loop failed at iteration {loop_count}: {type(e).__name__}: {e}")
            raise
        
        # Verify we processed everything
        if loop_count != len(temp_file_groups):
            self.logger.warning(f"Loop count mismatch! Expected {len(temp_file_groups)}, got {loop_count}")
        
        if self._should_log_memory("INFO"):
            self.console.print(f"[green]✅ Successfully consolidated {total_consolidations} datasets[/green]")
        
        # Final verification: check if temp files still exist
        final_check_groups = self._discover_temp_files()
        if final_check_groups:
            self.logger.warning(f"{len(final_check_groups)} temp file groups still exist after consolidation")
            for (provider, month), files in final_check_groups.items():
                self.logger.warning(f"Remaining temp files: {provider}/{month} - {len(files)} files")
    
    def _discover_temp_files(self) -> Dict[tuple, List[str]]:
        """Discover all temp files for this dataset, grouped by (provider, month)."""
        temp_file_groups = {}
        
        # List all files in the generate stage for this dataset using pagination
        prefix = f"czbillgen/datasets/{self.current_dataset_name}/generate/"
        
        # Use paginator to handle >1000 objects
        paginator = self.s3_service.s3_client.get_paginator('list_objects_v2')
        page_iterator = paginator.paginate(
            Bucket=self.s3_service.bucket,
            Prefix=prefix
        )
        
        total_objects = 0
        pages_processed = 0
        
        # Parse temp files and group by provider/month (paginated)
        temp_file_count = 0
        
        for page in page_iterator:
            pages_processed += 1
            
            if 'Contents' not in page:
                continue
            
            page_objects = len(page['Contents'])
            total_objects += page_objects
            
            for i, obj in enumerate(page['Contents']):
                key = obj['Key']
                
                # Parse key: czbillgen/datasets/{dataset}/generate/{provider}/{month}_temp_{timestamp}.parquet
                path_parts = key.split('/')
                if len(path_parts) >= 5:
                    provider = path_parts[-2]
                    filename = path_parts[-1]
                    
                    # Check if this is a temp file
                    if '_temp_' in filename and filename.endswith('.parquet'):
                        temp_file_count += 1
                        
                        # Extract original month from temp filename
                        month_part = filename.replace('.parquet', '').split('_temp_')[0]
                        
                        group_key = (provider, month_part)
                        if group_key not in temp_file_groups:
                            temp_file_groups[group_key] = []
                        temp_file_groups[group_key].append(key)
                        
        
        total_groups = len(temp_file_groups)
        
        # Verify temp file counting accuracy
        total_files_in_groups = sum(len(files) for files in temp_file_groups.values())
        
        if temp_file_count != total_files_in_groups:
            raise RuntimeError(f"Temp file count mismatch: found {temp_file_count} files but grouped {total_files_in_groups}")
        
        # Log summary if requested
        if self._should_log_memory("DEBUG"):
            self.logger.debug(f"S3 pagination processed {total_objects} objects across {pages_processed} pages, found {total_files_in_groups} temp files in {total_groups} groups")
        
        return temp_file_groups
    
    def _consolidate_provider_month(self, provider: str, month: str, temp_file_keys: List[str]):
        """Consolidate all temp files for a specific provider/month into final file."""
        
        if self._should_log_memory("DEBUG"):
            self.console.print(f"[dim]🔗 Consolidating {len(temp_file_keys)} temp files for {provider}/{month}[/dim]")
        
        # Read all temp files and collect DataFrames
        dataframes = []
        total_rows = 0
        
        for i, temp_key in enumerate(temp_file_keys, 1):
            # Read temp file using S3DataService helper (works around Polars S3 credential issues)
            temp_df = self.s3_service.read_parquet_from_s3(temp_key)
            dataframes.append(temp_df)
            total_rows += len(temp_df)
            
            if self._should_log_memory("DEBUG"):
                self.console.print(f"[dim]  📄 Read {len(temp_df):,} rows from {temp_key.split('/')[-1]}[/dim]")
        
        # Merge all DataFrames
        if len(dataframes) == 1:
            final_df = dataframes[0]
        else:
            final_df = pl.concat(dataframes, how="vertical")
        
        # Upload consolidated final file
        self.s3_service.upload_dataset(self.current_dataset_name, provider, month, final_df)
        
        if self._should_log_memory("INFO"):
            self.console.print(f"[green]✅ {provider}/{month}: consolidated {len(temp_file_keys)} files → {total_rows:,} rows[/green]")
        
        # Clean up temp files after successful consolidation
        self._cleanup_temp_files(temp_file_keys)
        return True  # Add explicit return to trace completion
    
    def _cleanup_temp_files(self, temp_file_keys: List[str]):
        """Clean up temporary files after successful consolidation."""
        
        if not temp_file_keys:
            return
        
        # Delete temp files in batch
        objects_to_delete = [{'Key': key} for key in temp_file_keys]
        
        try:
            response = self.s3_service.s3_client.delete_objects(
                Bucket=self.s3_service.bucket,
                Delete={'Objects': objects_to_delete}
            )
            
            # Check for errors in the response
            if 'Errors' in response and response['Errors']:
                self.logger.error(f"Some temp files could not be deleted: {response['Errors']}")
                
        except Exception as e:
            self.logger.error(f"Failed to delete temp files: {type(e).__name__}: {e}")
            # Don't raise - consolidation succeeded, cleanup is best-effort
        
        
        if self._should_log_memory("DEBUG"):
            self.console.print(f"[dim]🧹 Successfully cleaned up {len(temp_file_keys)} temp files[/dim]")
        
    
    def _collect_results(
        self,
        generation_results: Dict[str, Any],
        analysis_results: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Collect and format final results."""
        line_item_stats = self.line_item_tracker.get_summary_stats()
        provider_stats = self.line_item_tracker.get_all_provider_stats()
        
        return {
            "dataset_name": self.current_dataset_name,
            "generation": generation_results,
            "analysis": analysis_results,
            "line_items": {
                "summary": line_item_stats,
                "providers": {stats.provider: {
                    "total_line_items": stats.total_line_items,
                    "month_count": stats.month_count,
                    "min_month": stats.min_month,
                    "max_month": stats.max_month,
                    "avg_per_month": stats.avg_line_items_per_month
                } for stats in provider_stats},
                "tracking_table": self.line_item_tracker.format_summary_table()
            },
            "failed_jobs": self.failed_jobs,
            "success": len(self.failed_jobs) == 0
        }
    
    def _validate_data_chunk(self, chunk_df: pl.DataFrame, job: GenerationJob) -> None:
        """
        Validate data chunk immediately without accumulating in memory.
        
        Performs streaming validation on individual data chunks to detect
        issues early without memory overhead.
        
        Args:
            chunk_df: Data chunk to validate
            job: Generation job for context
            
        Raises:
            ValidationError: If chunk validation fails
        """
        if chunk_df.is_empty():
            return
        
        # Basic completeness validation
        completeness_result = validate_basic_completeness(
            chunk_df, job.start_date, job.end_date
        )
        if not completeness_result.success:
            raise ValidationError(f"Chunk validation failed: {completeness_result.message}")
        
        # Cost variance validation
        variance_result = validate_cost_variance(chunk_df)
        if not variance_result.success:
            raise ValidationError(f"Chunk cost variance failed: {variance_result.message}")
    
    def _check_memory_pressure_and_gc(self) -> None:
        """
        Check for memory pressure and force garbage collection if needed.
        
        Proactively manages memory usage to prevent out-of-memory errors
        during large dataset generation.
        """
        import psutil
        import gc
        import os
        
        # Get current process memory usage
        process = psutil.Process(os.getpid())
        memory_info = process.memory_info()
        memory_mb = memory_info.rss / 1024 / 1024
        
        # Check if memory usage exceeds 80% of target
        target_mb = self.target_memory_bytes / 1024 / 1024
        if memory_mb > target_mb * 0.8:
            # Force garbage collection
            gc.collect()
            
            # Evict cache entries if available
            if hasattr(self.reference_data_cache, 'evict_if_memory_pressure'):
                evicted = self.reference_data_cache.evict_if_memory_pressure(threshold_percentage=0.75)
                if evicted > 0 and self._should_log_memory("INFO"):
                    self.console.print(f"🔄 Memory pressure detected - evicted {evicted} cache entries")
            
            # Log memory usage only if configured
            if self._should_log_memory("DEBUG"):
                new_memory_info = process.memory_info()
                new_memory_mb = new_memory_info.rss / 1024 / 1024
                self.console.print(f"🔄 Memory: {memory_mb:.0f}MB → {new_memory_mb:.0f}MB (target: {target_mb:.0f}MB)")
    
    def _validate_job_data(self, df: pl.DataFrame, job: GenerationJob) -> None:
        """
        Validate single job data quality before S3 upload.
        
        Performs strict validation ensuring every job has:
        - Complete, valid data (no nulls, no zero costs)
        - Sufficient cost variance (no flat costs allowed)
        - Correct total cost (matches job.total_cost)
        
        Args:
            df: Generated DataFrame for this job
            job: Generation job configuration
            
        Raises:
            GenerationError: If any validation fails
        """
        # 1. Basic completeness validation
        completeness_result = validate_basic_completeness(
            df, job.start_date, job.end_date
        )
        if not completeness_result.success:
            raise GenerationError(
                f"Job '{job.name}' failed completeness validation: {completeness_result.message}"
            )
        
        # 2. Cost variance validation - no flat costs allowed
        variance_result = validate_cost_variance(df)
        if not variance_result.success:
            raise GenerationError(
                f"Job '{job.name}' failed variance validation: {variance_result.message}"
            )
        
        # 3. Cost conservation validation for this job's expected cost
        conservation_result = validate_cost_conservation(df, job.total_cost)
        if not conservation_result.success:
            raise GenerationError(
                f"Job '{job.name}' failed cost conservation: {conservation_result.message}"
            )
    
    def _write_success_marker(self, dataset_name: str) -> None:
        """
        Write _SUCCESS marker file to indicate complete dataset generation.
        
        This marker file allows data consumers to distinguish between
        complete and incomplete datasets.
        
        Args:
            dataset_name: Name of the generated dataset
        """
        from src.services.success_marker_service import SuccessMarkerService, create_processing_summary, create_validation_summary
        
        # Create standardized processing summary based on generation results
        total_jobs = len(self.all_generation_results) if hasattr(self, 'all_generation_results') else 0
        failed_jobs = len(self.failed_jobs)
        processed_jobs = total_jobs - failed_jobs
        
        processing_summary = create_processing_summary(
            total_items=total_jobs,
            processed_items=processed_jobs,
            failed_items=failed_jobs,
            additional_metrics={
                "line_items_generated": sum(result.get("line_items", 0) for result in self.all_generation_results) if hasattr(self, 'all_generation_results') else 0,
                "memory_usage_bytes": self.target_memory_bytes,
                "generation_engine": "streaming",
                "org_name": self.org_name
            }
        )
        
        # Create validation summary (Generate has validation)
        validation_summary = create_validation_summary(
            validation_performed=True,
            validation_passed=True,  # Only written if validation passed
            total_validations=3,  # Basic completeness, cost variance, cost conservation
            failed_validations=0,
            validation_details={
                "cost_conservation_validated": True,
                "basic_completeness_validated": True,
                "cost_variance_validated": True
            }
        )
        
        # Use standardized success marker service
        success_marker_service = SuccessMarkerService(self.s3_service)
        success = success_marker_service.create_success_marker(
            dataset_name=dataset_name,
            stage="generate",
            processing_summary=processing_summary,
            validation_summary=validation_summary,
            additional_metadata={
                "generator": "czbillgen-streaming-engine",
                "target_memory_bytes": self.target_memory_bytes,
                "bytes_per_record": self.bytes_per_record,
                "use_optimized_s3": self.use_optimized_s3,
                "use_vectorized_operations": self.use_vectorized_operations
            }
        )
        
        if success:
            self.console.print(f"[green]✓ Success marker written for dataset: {dataset_name}[/green]")
        else:
            self.console.print(f"[yellow]⚠️ Failed to write success marker[/yellow]")
            self.console.print(f"[yellow]Dataset '{dataset_name}' is complete but unmarked[/yellow]")
    
    def _get_current_phase_name(self, error: Exception) -> str:
        """
        Determine which phase failed based on the error and current state.
        
        Args:
            error: The exception that occurred
            
        Returns:
            Human-readable phase name for error reporting
        """
        # Use error messages and context to determine phase
        error_msg = str(error).lower()
        
        if "reference" in error_msg or "filter" in error_msg:
            return "Phase 1 (Job Analysis)"
        elif "memory" in error_msg or "split" in error_msg:
            return "Phase 2 (Memory Optimization)"  
        elif "generation" in error_msg or "resource" in error_msg:
            return "Phase 3 (Data Generation)"
        elif "collect" in error_msg or "result" in error_msg:
            return "Phase 4 (Results Collection)"
        elif "finalize" in error_msg or "temp" in error_msg:
            return "Phase 5 (Finalization)"
        elif "success" in error_msg or "marker" in error_msg:
            return "Phase 6 (Success Marker)"
        else:
            return "Unknown Phase"
    
    def _emergency_cleanup(self, dataset_name: str) -> None:
        """
        Emergency cleanup on generation failure to prevent stealth error state.
        
        Removes any temp files and partial data to ensure the system fails cleanly
        rather than leaving orphaned temp files that create confusion.
        
        Args:
            dataset_name: Name of the failed dataset
        """
        self.console.print("[yellow]🧹 Cleaning up temp files to prevent stealth error state...[/yellow]")
        
        # Clear any accumulated data in memory
        if self.accumulated_data:
            memory_mb = self.current_accumulation_size / (1024 * 1024)
            self.console.print(f"[yellow]   • Clearing {len(self.accumulated_data)} accumulated datasets ({memory_mb:.1f}MB)[/yellow]")
            self.accumulated_data.clear()
            self.current_accumulation_size = 0
        
        # Check for and remove any temp files that were already uploaded
        temp_file_groups = self._discover_temp_files()
        if temp_file_groups:
            self._consolidate_temp_files(temp_file_groups)
            
            total_temp_files = sum(len(files) for files in temp_file_groups.values())
            self.console.print(f"[yellow]   • Found {total_temp_files} temp files for {len(temp_file_groups)} provider/month combinations[/yellow]")
            
            # Delete all temp files to prevent stealth state
            for (provider, month), temp_file_keys in temp_file_groups.items():
                self.console.print(f"[yellow]   • Removing temp files for {provider}/{month}[/yellow]")
                self._cleanup_temp_files(temp_file_keys)
            
            self.console.print(f"[green]✅ Cleaned up {total_temp_files} temp files[/green]")
        else:
            self.console.print("[green]   • No temp files found to clean up[/green]")