"""
DataManager service for optimal data operations.

Provides data operation services USED BY pipeline, not controlling pipeline execution.
Focuses on memory-efficient loading/saving, automatic optimization, and performance monitoring.
"""

import asyncio
import time
from collections.abc import Callable
from contextlib import contextmanager
from pathlib import Path
from typing import (
    Any,
    Dict,
    List,
    Optional,
    Tuple,
)

import polars as pl

from .config import DataManagerConfig
from .context import PerformanceContext
from .models import PerformanceMetrics
from .monitoring import PerformanceMonitor
from .storage import SimpleStorageManager
from .strategies import ProcessingStrategySelector


class DataManager:
    """Service for optimal data operations - USED BY pipeline, not controlling it"""

    def __init__(self, run_dir: Path, config: Optional[DataManagerConfig] = None):
        self.run_dir = run_dir
        self.config = config or DataManagerConfig.auto_detect()

        self.strategy_selector = ProcessingStrategySelector(self.config)
        self.performance_monitor = PerformanceMonitor(self.config.memory_limit_gb)

        # Simple storage management (extracted valuable pieces from Phase 4)
        self.storage_manager = SimpleStorageManager(run_dir)

        # Track current operation for metrics
        self.current_strategy: Optional[str] = None
        self.last_operation_time: float = 0.0

    # Service methods called BY pipeline
    def load_optimally(self, source: Path) -> pl.DataFrame:
        """Load data with optimal strategy based on size and memory constraints"""
        # Validate source file exists
        self._validate_source_file_exists(source)

        # Select and apply optimal loading strategy
        return self._apply_optimal_loading_strategy(source)

    def save_optimally(self, data: pl.DataFrame, stage_name: str, merge_mode: str = "overwrite") -> Path:
        """
        Save data with optimal format, compression, and partitioning.

        Args:
            data: DataFrame to save
            stage_name: Stage name for output directory path
            merge_mode: How to handle existing files ("overwrite" or "append")

        Returns:
            Path to output directory
        """
        # Validate input parameters
        self._validate_save_parameters(data, stage_name, merge_mode)

        # Prepare output directory
        output_dir = self._prepare_output_directory(stage_name)

        # Determine and apply optimal saving strategy
        return self._apply_optimal_saving_strategy(data, output_dir, merge_mode)

    def get_optimal_strategy(self, source: Path) -> str:
        """Determine optimal processing strategy without executing"""
        strategy = self.strategy_selector.select_strategy(source, "analyze")
        return strategy.name

    @contextmanager
    def monitor_performance(self, operation_name: str):
        """Context manager for performance monitoring"""
        context = PerformanceContext(self.performance_monitor, operation_name)
        try:
            with context:
                yield context
        finally:
            self.last_operation_time = context.elapsed_time

    def get_performance_metrics(self) -> PerformanceMetrics:
        """Get current performance metrics"""
        return PerformanceMetrics(
            processing_time_seconds=self.last_operation_time,
            peak_memory_usage_gb=self.performance_monitor.get_peak_memory(),
            optimization_strategy=self.current_strategy or "none",
            memory_efficiency_score=self._calculate_memory_efficiency(),
            storage_compression_ratio=self._calculate_compression_ratio(),
        )

    def get_performance_summary(self) -> dict:
        """Get comprehensive performance summary (legacy method)"""
        metrics = self.get_performance_metrics()
        return {
            "total_time_seconds": metrics.processing_time_seconds,
            "peak_memory_gb": metrics.peak_memory_usage_gb,
            "optimization_summary": f"Used {metrics.optimization_strategy} strategy",
            "memory_efficiency": metrics.memory_efficiency_score,
        }

    def cleanup_storage(self, keep_runs: int = 5) -> dict:
        """Simple storage cleanup with smart defaults"""
        return self.storage_manager.cleanup_old_runs(keep_runs)

    def get_storage_info(self) -> dict:
        """Basic storage information"""
        return self.storage_manager.get_storage_usage()

    def load_stage_data(self, stage_name: str) -> pl.DataFrame:
        """
        Load data for a specific stage, handling partitioned data automatically.

        This method handles the case where DataManager saves data in partitioned format
        (provider-specific files) and loads/combines them into a unified DataFrame.

        Args:
            stage_name: Name of the stage to load data for

        Returns:
            Unified DataFrame with all data for the stage
        """
        stage_dir = self.run_dir / stage_name

        # First check for standard data.parquet file
        standard_file = stage_dir / "data.parquet"
        if standard_file.exists():
            df = pl.read_parquet(standard_file)
            # Check if this is summary data vs actual data
            if self._is_summary_data(df):
                # This is summary data, try to load actual partitioned data
                return self._load_partitioned_data(stage_dir)
            return df

        # Try to load partitioned data directly
        return self._load_partitioned_data(stage_dir)

    def _is_summary_data(self, df: pl.DataFrame) -> bool:
        """
        Detect if DataFrame contains summary statistics rather than actual billing data.

        Summary data characteristics:
        - Very few rows (typically 1)
        - Contains metadata columns like 'stage', 'total_files_created'
        - Missing actual billing columns like service-specific columns
        """
        if df.height > 10:  # Summary data is typically 1-2 rows
            return False

        summary_indicators = {
            "stage",
            "total_files_created",
            "total_rows",
            "processing_time",
        }
        billing_indicators = {"service", "cost", "cloud_provider"}

        has_summary_cols = any(col in df.columns for col in summary_indicators)
        has_billing_cols = any(col in df.columns for col in billing_indicators)

        # If it has summary columns but lacks billing columns, it's likely summary data
        return has_summary_cols and not has_billing_cols

    def _load_partitioned_data(self, stage_dir: Path) -> pl.DataFrame:
        """
        Load and combine partitioned data files from a stage directory.

        Handles DataManager's partitioned storage strategy where data is split
        by cloud provider (AWS.parquet, Azure.parquet, GCP.parquet).
        """
        # Look for provider-specific files
        provider_files = []
        for provider in ["AWS", "Azure", "GCP"]:
            provider_file = stage_dir / f"{provider}.parquet"
            if provider_file.exists():
                provider_files.append(provider_file)

        if not provider_files:
            # No partitioned files found, check for month-based partitioning (generate stage)
            return self._load_month_partitioned_data(stage_dir)

        # Load and combine provider files with schema alignment
        return self._combine_files_with_schema_alignment(provider_files)

    def _load_month_partitioned_data(self, stage_dir: Path) -> pl.DataFrame:
        """
        Load month-partitioned data (generate stage pattern: Provider/YYYY-MM/data.parquet).
        """
        data_files = []

        # Look for provider/month structure
        for provider_dir in stage_dir.iterdir():
            if provider_dir.is_dir() and provider_dir.name in ["AWS", "Azure", "GCP"]:
                for month_dir in provider_dir.iterdir():
                    if month_dir.is_dir():
                        data_file = month_dir / "data.parquet"
                        if data_file.exists():
                            data_files.append(data_file)

        if not data_files:
            raise FileNotFoundError(f"No data files found in {stage_dir}")

        return self._combine_files_with_schema_alignment(data_files)

    def _combine_files_with_schema_alignment(self, file_paths: List[Path]) -> pl.DataFrame:
        """
        Load and combine multiple parquet files with automatic schema alignment.

        Handles cases where files may have slightly different schemas by:
        1. Collecting all unique columns across files
        2. Loading each file with consistent schema (adding missing columns as null)
        3. Combining into unified DataFrame
        """
        if not file_paths:
            raise ValueError("No files provided to combine")

        # First pass: collect all columns across all files
        all_columns = set()
        file_schemas = {}

        for file_path in file_paths:
            try:
                # Read just the schema without loading data
                lazy_frame = pl.scan_parquet(file_path)
                df_schema = lazy_frame.schema
                columns = set(df_schema.keys())
                all_columns.update(columns)
                file_schemas[file_path] = df_schema
            except Exception as e:
                print(f"Warning: Could not read schema from {file_path}: {e}")
                continue

        if not all_columns:
            raise ValueError("Could not read schemas from any provided files")

        all_columns = sorted(all_columns)  # Consistent ordering

        # Second pass: load data with consistent schema
        dataframes = []
        for file_path in file_paths:
            try:
                df = pl.read_parquet(file_path)

                # Add missing columns as null
                for col in all_columns:
                    if col not in df.columns:
                        # Determine appropriate null type based on similar columns from other files
                        null_type = self._infer_column_type(col, file_schemas)
                        df = df.with_columns(pl.lit(None, dtype=null_type).alias(col))

                # Ensure column order is consistent
                df = df.select(all_columns)
                dataframes.append(df)

            except Exception as e:
                print(f"Warning: Could not load data from {file_path}: {e}")
                continue

        if not dataframes:
            raise ValueError("Could not load data from any provided files")

        # Combine all dataframes
        return pl.concat(dataframes, how="vertical")

    def _infer_column_type(self, column_name: str, schemas: dict) -> pl.DataType:
        """
        Infer appropriate data type for a missing column based on other schemas.
        """
        # Look for the column in other schemas
        for schema in schemas.values():
            if column_name in schema:
                return schema[column_name]

        # Default fallback based on column name patterns
        if "cost" in column_name.lower() or "price" in column_name.lower():
            return pl.Float64
        elif "count" in column_name.lower() or "id" in column_name.lower():
            return pl.Int64
        elif "date" in column_name.lower() or "time" in column_name.lower():
            return pl.Utf8  # Will be parsed as needed
        else:
            return pl.Utf8  # Default to string

    def _calculate_memory_efficiency(self) -> float:
        """Calculate memory efficiency score (0-100)"""
        # Simple heuristic: efficiency based on peak memory vs limit
        if self.config.memory_limit_gb == 0:
            return 100.0

        peak_memory = self.performance_monitor.get_peak_memory()
        efficiency = max(0, 100 - (peak_memory / self.config.memory_limit_gb * 100))
        return min(100.0, efficiency)

    def _calculate_compression_ratio(self) -> Optional[float]:
        """Calculate storage compression ratio if available"""
        # This would be enhanced to track actual compression in future
        if self.config.compression_enabled:
            return 0.7  # Approximate compression ratio
        return None

    def _is_large_dataset(self, data: pl.DataFrame) -> bool:
        """Determine if dataset should be partitioned"""
        # Simple heuristic: partition if > 1M rows or estimated size > chunk_size
        estimated_size_mb = data.height * 100 / (1024 * 1024)  # ~100 bytes per row
        return data.height > 1_000_000 or estimated_size_mb > self.config.chunk_size_mb

    def _save_partitioned(self, data: pl.DataFrame, output_dir: Path, compression: Optional[str]):
        """Save large dataset with partitioning (legacy method for backward compatibility)"""
        # Use new merge-aware method with overwrite mode
        self._save_partitioned_with_merge(data, output_dir, compression, "overwrite")

    def _save_single_file_with_merge(
        self,
        data: pl.DataFrame,
        output_file: Path,
        compression: Optional[str],
        merge_mode: str,
    ):
        """Save single file with merge support (append or overwrite)"""
        if merge_mode == "overwrite" or not output_file.exists():
            # Standard overwrite or new file
            data.write_parquet(output_file, compression=compression)
        else:
            # Append mode: merge with existing data
            try:
                existing_data = pl.read_parquet(output_file)
                merged_data = pl.concat([existing_data, data], how="vertical")

                # Atomic write: write to temp file then move
                temp_file = output_file.with_suffix(".tmp.parquet")
                merged_data.write_parquet(temp_file, compression=compression)
                temp_file.replace(output_file)

            except Exception as e:
                # Fallback to overwrite if merge fails
                print(f"Warning: Failed to append to {output_file}, overwriting instead: {e}")
                data.write_parquet(output_file, compression=compression)

    def _save_partitioned_with_merge(
        self,
        data: pl.DataFrame,
        output_dir: Path,
        compression: Optional[str],
        merge_mode: str,
    ):
        """Save partitioned data with merge support"""
        if "cloud_provider" in data.columns:
            providers = data["cloud_provider"].unique().to_list()
            for provider in providers:
                provider_data = data.filter(pl.col("cloud_provider") == provider)
                provider_file = output_dir / f"{provider}.parquet"

                if merge_mode == "overwrite" or not provider_file.exists():
                    # Standard overwrite or new file
                    provider_data.write_parquet(provider_file, compression=compression)
                else:
                    # Append mode: merge with existing provider data
                    try:
                        existing_data = pl.read_parquet(provider_file)
                        merged_data = pl.concat([existing_data, provider_data], how="vertical")

                        # Atomic write: write to temp file then move
                        temp_file = provider_file.with_suffix(".tmp.parquet")
                        merged_data.write_parquet(temp_file, compression=compression)
                        temp_file.replace(provider_file)

                    except Exception as e:
                        # Fallback to overwrite if merge fails
                        print(f"Warning: Failed to append to {provider_file}, overwriting instead: {e}")
                        provider_data.write_parquet(provider_file, compression=compression)
        else:
            # Fallback to single file if no cloud_provider column
            output_file = output_dir / "data.parquet"
            self._save_single_file_with_merge(data, output_file, compression, merge_mode)

    def _calculate_data_size(self, location: Path) -> float:
        """Calculate data size in GB"""
        total_size = 0
        try:
            if location.is_file():
                total_size = location.stat().st_size
            else:
                for file_path in location.rglob("*"):
                    if file_path.is_file():
                        total_size += file_path.stat().st_size
        except (OSError, PermissionError):
            pass
        return total_size / (1024**3)

    # V1 File-by-File Processing Methods (Phase 5.1)
    def get_provider_months(self, stage_name: str) -> List[Tuple[str, str]]:
        """
        Get list of (provider, month) tuples for file-by-file processing.

        Scans stage directory for provider/month structure and returns all available
        combinations for V1-style file-by-file processing.

        Args:
            stage_name: Name of the stage (e.g., "generate", "k8s")

        Returns:
            List of (provider, month) tuples, e.g., [("AWS", "2024-01"), ("GCP", "2024-01")]
        """
        stage_dir = self.run_dir / stage_name

        if not stage_dir.exists():
            return []

        provider_months = []

        # Look for provider/month directory structure (generate stage pattern)
        for provider_dir in stage_dir.iterdir():
            if provider_dir.is_dir():
                # Check for parquet files in provider/month/data.parquet structure
                for month_dir in provider_dir.iterdir():
                    if month_dir.is_dir() and self._is_valid_month_format(month_dir.name):
                        # Check if there's actually data in this month
                        data_file = month_dir / "data.parquet"
                        if data_file.exists():
                            provider_months.append((provider_dir.name, month_dir.name))

        # All stages must use consistent Provider/YYYY-MM/data.parquet structure

        # Sort for consistent processing order
        return sorted(provider_months)

    def load_stage_data_chunk(self, stage_name: str, provider: str, month: str) -> pl.DataFrame:
        """
        Load data for a specific provider/month chunk.

        This method loads individual provider/month files to enable V1-style
        file-by-file processing with constant memory usage (~513K rows max).

        Args:
            stage_name: Name of the stage
            provider: Provider name (e.g., "AWS", "GCP")
            month: Month in YYYY-MM format (e.g., "2024-01")

        Returns:
            DataFrame with data for the specific provider/month
        """
        stage_dir = self.run_dir / stage_name

        # All stages must use consistent Provider/YYYY-MM/data.parquet structure
        provider_month_file = stage_dir / provider / month / "data.parquet"
        if provider_month_file.exists():
            return pl.read_parquet(provider_month_file)

        raise FileNotFoundError(f"No data found for {provider}/{month} in stage '{stage_name}'")

    def save_stage_data_chunk(self, stage_name: str, data: pl.DataFrame, provider: str, month: str) -> Path:
        """
        Save data for a specific provider/month chunk.

        Saves processed data in provider/month structure for efficient file-by-file
        processing pattern.

        Args:
            stage_name: Name of the stage
            data: Processed DataFrame to save
            provider: Provider name
            month: Month in YYYY-MM format

        Returns:
            Path to the saved file
        """
        # Validate input
        if data is None or data.is_empty():
            raise ValueError(f"Cannot save empty data for {provider}/{month}")

        # Create provider/month directory structure
        output_dir = self.run_dir / stage_name / provider / month
        output_dir.mkdir(parents=True, exist_ok=True)

        # Save with optimal compression
        output_file = output_dir / "data.parquet"
        compression = "snappy" if self.config.processing_preference == "speed" else "lz4"
        data.write_parquet(
            output_file,
            compression=compression if self.config.compression_enabled else None,
        )

        return output_file

    def process_stage_by_files(
        self,
        stage_func: Callable,
        input_stage: str,
        output_stage: str,
        config: Optional[Dict[str, Any]] = None,
        summary_log: Optional[Any] = None,
    ) -> Dict[str, Any]:
        """
        Process stage using V1 file-by-file pattern for memory efficiency.

        This method orchestrates file-by-file processing to achieve constant memory usage
        regardless of total dataset size, restoring V1's proven scalability pattern.

        Args:
            stage_func: Stage function to apply to each chunk
            input_stage: Name of input stage to read from
            output_stage: Name of output stage to write to
            config: Optional configuration for stage function
            summary_log: Optional summary log for tracking

        Returns:
            Processing summary with metrics and results
        """
        # Initialize processing context
        provider_months, processing_results = self._initialize_file_processing_context(input_stage, output_stage)

        # Process each provider/month file
        processed_chunks, total_rows_processed, total_rows_output, start_time = self._process_all_provider_month_files(
            provider_months, stage_func, input_stage, output_stage, config, summary_log, processing_results
        )

        # Finalize and return results
        return self._finalize_file_processing_results(
            processing_results,
            processed_chunks,
            total_rows_processed,
            total_rows_output,
            start_time,
            len(provider_months),
        )

    def _is_valid_month_format(self, month_str: str) -> bool:
        """Check if string is valid YYYY-MM month format"""
        try:
            if len(month_str) != 7 or month_str[4] != "-":
                return False
            year = int(month_str[:4])
            month = int(month_str[5:7])
            return 2020 <= year <= 2030 and 1 <= month <= 12
        except (ValueError, IndexError):
            return False

    def write_payload_chunks(
        self,
        data: pl.DataFrame,
        output_dir: Path,
        max_size_mb: float = 4.5,
        provider: str = None,
        month: str = None,
    ) -> List[Path]:
        """
        Write DataFrame as chunked JSON payload files with memory efficiency.

        Splits large datasets into multiple payload files that respect CloudZero API
        size limits while maintaining memory efficiency through DataManager patterns.

        Args:
            data: DataFrame to write as payload chunks
            output_dir: Directory to write payload files
            max_size_mb: Maximum size per payload file in MB
            provider: Provider name for filename organization (optional)
            month: Month for filename organization (optional)

        Returns:
            List of paths to created payload files
        """

        if data is None or data.is_empty():
            return []

        # Create payloads directory
        payloads_dir = output_dir / "payloads"
        payloads_dir.mkdir(parents=True, exist_ok=True)

        # Split data into 5MB chunks (legacy chunking logic)
        # Note: Upload functionality now uses S3-based approach
        max_chunk_size = 5 * 1024 * 1024  # 5MB in bytes
        
        # Simple chunking by row count (approximation)
        rows_per_chunk = max(1000, max_chunk_size // (len(data.columns) * 100))  # rough estimate
        
        payload_groups = []
        for i in range(0, len(data), rows_per_chunk):
            chunk = data.slice(i, rows_per_chunk)
            payload_groups.append({
                "data": chunk,
                "provider": "default",
                "connection": "default"
            })

        payload_files = []
        for group in payload_groups:
            # Create filename with provider/month if available
            if provider and month:
                filename = f"{provider}_{month}_payload_{group.payload_index}.jsonl"
            else:
                filename = f"payload_{group.payload_index}.jsonl"

            payload_file = payloads_dir / filename

            # Handle column dropping logic (copied from persister.py)
            if "billing_month" in group.dataframe.columns:
                df_to_write = group.dataframe.drop("billing_month", "cloud_provider")
            else:
                df_to_write = group.dataframe

            # Write payload file efficiently using async I/O with standard library
            asyncio.run(self._write_ndjson_async(df_to_write, payload_file))

            payload_files.append(payload_file)

        return payload_files

    async def _write_ndjson_async(self, df: pl.DataFrame, file_path: Path) -> None:
        """Write DataFrame as JSONL using async I/O with standard library."""
        json_data = df.write_ndjson()
        loop = asyncio.get_event_loop()
        # Use thread executor for file I/O to avoid blocking
        await loop.run_in_executor(None, lambda: file_path.write_text(json_data))

    def discover_payload_files(self, payload_dir: Path) -> List[Tuple[str, str, Path]]:
        """
        Discover payload files with provider/month organization.

        Extends existing DataManager discovery patterns to payload files.

        Args:
            payload_dir: Directory containing payload files

        Returns:
            List of (provider, month, file_path) tuples
        """
        if not payload_dir.exists():
            return []

        payload_files = []
        for jsonl_file in payload_dir.glob("*.jsonl"):
            # Parse filename patterns like: AWS_2024-04_payload_0.jsonl
            parts = jsonl_file.stem.split("_")
            if len(parts) >= 4 and parts[2] == "payload":
                provider = parts[0].upper()
                month = parts[1]
                if self._is_valid_month_format(month):
                    payload_files.append((provider, month, jsonl_file))

        # Sort by provider, then month for consistent processing
        return sorted(payload_files)

    # Helper functions for load_optimally
    def _validate_source_file_exists(self, source: Path) -> None:
        """Validate that the source file exists and provide helpful error messages."""
        if not source.exists():
            from pathlib import Path as PathLib

            current_dir = PathLib.cwd()
            absolute_path = source.resolve()

            error_msg = f"Data source not found: {source}\n"
            error_msg += f"Tried absolute path: {absolute_path}\n"
            error_msg += f"Current working directory: {current_dir}\n"

            # Check if user might need to be in cloudzero-demo-billing directory
            expected_dir = current_dir / "cloudzero-demo-billing"
            if expected_dir.exists() and current_dir.name != "cloudzero-demo-billing":
                error_msg += (
                    "💡 Tip: Run 'cd cloudzero-demo-billing' first, or provide full path to your reference data file"
                )
            elif not (current_dir / "orgs").exists():
                error_msg += "💡 Tip: Make sure you're running from the cloudzero-demo-billing directory"

            raise FileNotFoundError(error_msg)

    def _apply_optimal_loading_strategy(self, source: Path) -> pl.DataFrame:
        """Select and apply optimal loading strategy for the source file."""
        # Select optimal strategy
        strategy = self.strategy_selector.select_strategy(source, "load")
        self.current_strategy = strategy.name

        # Load data using strategy
        return strategy.load_data(source)

    # Helper functions for save_optimally
    def _validate_save_parameters(self, data: pl.DataFrame, stage_name: str, merge_mode: str) -> None:
        """Validate input parameters for save operations."""
        if data is None:
            raise ValueError(f"Cannot save None DataFrame for stage '{stage_name}'")

        if not isinstance(data, pl.DataFrame):
            raise TypeError(f"Expected pl.DataFrame for stage '{stage_name}', got {type(data)}")

        if data.is_empty():
            raise ValueError(f"Cannot save empty DataFrame for stage '{stage_name}'")

        if merge_mode not in ("overwrite", "append"):
            raise ValueError(f"merge_mode must be 'overwrite' or 'append', got '{merge_mode}'")

    def _prepare_output_directory(self, stage_name: str) -> Path:
        """Prepare and create the output directory for the stage."""
        output_dir = self.run_dir / stage_name
        output_dir.mkdir(parents=True, exist_ok=True)
        return output_dir

    def _apply_optimal_saving_strategy(self, data: pl.DataFrame, output_dir: Path, merge_mode: str) -> Path:
        """Determine and apply optimal saving strategy for the data."""
        # Determine optimal strategy for saving
        if self._is_large_dataset(data):
            self.current_strategy = "chunked"  # Large datasets use chunked strategy
        else:
            self.current_strategy = "in_memory"  # Small datasets use in_memory strategy

        # Determine compression settings
        compression = self._determine_compression_settings()

        # Apply partitioning strategy based on dataset size
        if self._is_large_dataset(data):
            self._save_partitioned_with_merge(data, output_dir, compression, merge_mode)
        else:
            output_file = output_dir / "data.parquet"
            self._save_single_file_with_merge(data, output_file, compression, merge_mode)

        return output_dir

    def _determine_compression_settings(self) -> Optional[str]:
        """Determine optimal compression settings based on configuration."""
        if self.config.compression_enabled:
            return "snappy" if self.config.processing_preference == "speed" else "lz4"
        return None

    # Helper functions for process_stage_by_files
    def _initialize_file_processing_context(
        self, input_stage: str, output_stage: str
    ) -> Tuple[List[Tuple[str, str]], Dict[str, Any]]:
        """Initialize processing context for file-by-file processing."""
        provider_months = self.get_provider_months(input_stage)

        if not provider_months:
            raise ValueError(f"No provider/month data found in stage '{input_stage}'")

        total_chunks = len(provider_months)
        print(f"🔄 Processing {total_chunks} provider/month files for {output_stage} stage...")

        processing_results = {
            "provider_months_processed": [],
            "total_chunks": total_chunks,
            "chunks_processed": 0,
            "total_input_rows": 0,
            "total_output_rows": 0,
            "processing_time_seconds": 0.0,
            "memory_efficiency_achieved": True,
        }

        return provider_months, processing_results

    def _process_all_provider_month_files(
        self,
        provider_months: List[Tuple[str, str]],
        stage_func: Callable,
        input_stage: str,
        output_stage: str,
        config: Optional[Dict[str, Any]],
        summary_log: Optional[Any],
        processing_results: Dict[str, Any],
    ) -> Tuple[int, int, int, float]:
        """Process all provider/month files and track metrics."""
        processed_chunks = 0
        total_rows_processed = 0
        total_rows_output = 0
        start_time = time.time()
        total_chunks = len(provider_months)

        for provider, month in provider_months:
            try:
                chunk_result = self._process_single_provider_month_file(
                    provider, month, stage_func, input_stage, output_stage, config, summary_log, processing_results
                )

                input_rows, output_rows = chunk_result
                total_rows_processed += input_rows
                total_rows_output += output_rows
                processed_chunks += 1

                # Progress reporting
                self._report_processing_progress(processed_chunks, total_chunks)

            except Exception as e:
                # Handle processing errors
                if self._should_fail_fast_on_error(e, provider, month):
                    raise e
                # Continue with other chunks for non-critical errors
                continue

        return processed_chunks, total_rows_processed, total_rows_output, start_time

    def _process_single_provider_month_file(
        self,
        provider: str,
        month: str,
        stage_func: Callable,
        input_stage: str,
        output_stage: str,
        config: Optional[Dict[str, Any]],
        summary_log: Optional[Any],
        processing_results: Dict[str, Any],
    ) -> Tuple[int, int]:
        """Process a single provider/month file and return row counts."""
        # Load chunk data (V1 pattern: ~513K rows max)
        chunk_data = self.load_stage_data_chunk(input_stage, provider, month)
        input_rows = chunk_data.height

        # Apply stage function to chunk
        processed_chunk = self._apply_stage_function_to_chunk(stage_func, chunk_data, config, summary_log)

        output_rows = processed_chunk.height if processed_chunk is not None else 0

        if processed_chunk is not None and not processed_chunk.is_empty():
            # Save processed chunk
            output_path = self.save_stage_data_chunk(output_stage, processed_chunk, provider, month)

            # Track processing results
            processing_results["provider_months_processed"].append(
                {
                    "provider": provider,
                    "month": month,
                    "input_rows": input_rows,
                    "output_rows": output_rows,
                    "output_path": str(output_path),
                }
            )

        return input_rows, output_rows

    def _apply_stage_function_to_chunk(
        self, stage_func: Callable, chunk_data: pl.DataFrame, config: Optional[Dict[str, Any]], summary_log: Optional[Any]
    ) -> pl.DataFrame:
        """Apply stage function to chunk with appropriate parameters."""
        if config is not None and summary_log is not None:
            return stage_func(chunk_data, config, summary_log)
        elif config is not None:
            return stage_func(chunk_data, config)
        else:
            return stage_func(chunk_data)

    def _report_processing_progress(self, processed_chunks: int, total_chunks: int) -> None:
        """Report processing progress to user."""
        if total_chunks > 10 and processed_chunks % max(1, total_chunks // 10) == 0:
            progress_pct = (processed_chunks / total_chunks) * 100
            print(f"   ⏳ Progress: {processed_chunks}/{total_chunks} files ({progress_pct:.0f}%)")

    def _should_fail_fast_on_error(self, error: Exception, provider: str, month: str) -> bool:
        """Determine if error should cause immediate failure or be skipped."""
        # Simple error classification - validation errors contain "validation" in the message
        error_message = str(error).lower()

        if "validation" in error_message:
            # Validation errors should fail the entire stage (fail-fast)
            print(f"   ❌ Validation failed for {provider}/{month}: {error}")
            return True
        else:
            # Other processing errors can be logged and skipped
            print(f"   ⚠️  Error processing {provider}/{month}: {error}")
            return False

    def _finalize_file_processing_results(
        self,
        processing_results: Dict[str, Any],
        processed_chunks: int,
        total_rows_processed: int,
        total_rows_output: int,
        start_time: float,
        total_chunks: int,
    ) -> Dict[str, Any]:
        """Finalize processing results and print summary."""
        processing_time = time.time() - start_time

        # Update final results
        processing_results.update(
            {
                "chunks_processed": processed_chunks,
                "total_input_rows": total_rows_processed,
                "total_output_rows": total_rows_output,
                "processing_time_seconds": processing_time,
            }
        )

        # Print completion summary
        if processed_chunks == total_chunks:
            print(
                f"✅ Completed processing: {processed_chunks}/{total_chunks} files, {total_rows_processed:,} → {total_rows_output:,} rows in {processing_time:.1f}s"
            )
        else:
            print(
                f"⚠️  Partial processing completed: {processed_chunks}/{total_chunks} files, {total_rows_processed:,} → {total_rows_output:,} rows in {processing_time:.1f}s"
            )

        return processing_results
