"""
Stage adapters for worker.

Adapt existing stage engines to worker interface.
Each adapter converts job config to engine format and extracts result metadata.
"""

import os
import time
from pathlib import Path
from typing import Any, Protocol
import polars as pl
import boto3

# Import existing stage engines
from src.k8s.services import MonthProcessor
from src.k8s.config import K8sConfig


def get_per_month_statistics(dataset_name: str, stage: str) -> dict[str, Any]:
    """
    Extract per-month line item counts from S3 output files.

    Reads parquet files from S3 and returns statistics per provider-month.
    Highlights peak month for CloudZero platform billing awareness.

    Args:
        dataset_name: Dataset identifier
        stage: Stage name (generate, k8s, metadata, usage, discount)

    Returns:
        Dict with keys:
            - per_month: Dict[str, Dict] with provider-month stats
            - peak_month: Dict with peak month details
            - total_records: Total across all months
    """
    try:
        s3_bucket = os.environ.get("CZBILLGEN_S3_BUCKET")
        if not s3_bucket:
            return {"per_month": {}, "peak_month": {}, "total_records": 0}

        s3_client = boto3.client("s3")
        s3_prefix = f"czbillgen/datasets/{dataset_name}/{stage}/"

        # List all parquet files
        response = s3_client.list_objects_v2(
            Bucket=s3_bucket, Prefix=s3_prefix, Delimiter="/"
        )

        per_month_stats = {}
        total_records = 0

        # Get common prefixes (provider directories)
        for prefix_info in response.get("CommonPrefixes", []):
            provider_prefix = prefix_info["Prefix"]
            provider = provider_prefix.split("/")[-2]

            # List month files for this provider
            month_response = s3_client.list_objects_v2(
                Bucket=s3_bucket, Prefix=provider_prefix
            )

            for obj in month_response.get("Contents", []):
                key = obj["Key"]
                if not key.endswith(".parquet"):
                    continue

                # Extract month from filename (e.g., 2025-03.parquet)
                filename = key.split("/")[-1]
                month = filename.replace(".parquet", "")

                # Download and count records
                local_path = f"/tmp/{provider}_{month}.parquet"
                s3_client.download_file(s3_bucket, key, local_path)

                df = pl.read_parquet(local_path)
                record_count = len(df)
                total_records += record_count

                # Store stats
                per_month_stats[f"{provider}-{month}"] = {
                    "provider": provider,
                    "month": month,
                    "record_count": record_count,
                    "file_size_mb": round(obj["Size"] / (1024 * 1024), 2),
                }

                # Clean up
                os.remove(local_path)

        # Find peak month
        peak_month = {}
        if per_month_stats:
            peak_key = max(per_month_stats.keys(), key=lambda k: per_month_stats[k]["record_count"])
            peak_month = per_month_stats[peak_key].copy()
            peak_month["key"] = peak_key

        return {
            "per_month": per_month_stats,
            "peak_month": peak_month,
            "total_records": total_records,
        }

    except Exception as e:
        print(f"WARNING: Failed to extract per-month statistics: {e}")
        return {"per_month": {}, "peak_month": {}, "total_records": 0}


class NoOpConsole:
    """
    No-op console for headless worker execution.

    Provides Rich Console interface but discards all output.
    Used when running stages in worker without terminal output.
    """

    def print(self, *args, **kwargs):
        """Discard print output."""
        pass

    def log(self, *args, **kwargs):
        """Discard log output."""
        pass

    def status(self, *args, **kwargs):
        """Return no-op context manager."""
        return self

    def get_time(self):
        """Return current time (for progress tracking)."""
        import time
        return time.time()

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, *args):
        """Context manager exit."""
        pass

    def Progress(self, *args, **kwargs):
        """
        Return no-op progress context manager.

        Used by streaming_engine.py for progress bars in headless mode.
        Returns self to act as context manager with add_task/update methods.
        """
        return self

    def add_task(self, description: str, total: int | None = None) -> str:
        """
        No-op task creation - returns dummy task ID.

        Args:
            description: Task description (ignored)
            total: Total steps (ignored)

        Returns:
            Dummy task ID string
        """
        return "noop-task"

    def update(self, task_id: str | None = None, **kwargs) -> None:
        """
        No-op progress update.

        Args:
            task_id: Task ID to update (ignored)
            **kwargs: Progress parameters (ignored)
        """
        pass

    def __getattr__(self, name):
        """Catch all other attribute access."""
        # Return self for chaining (console.style.method())
        if name in ('style', 'highlighter', 'theme'):
            return self
        # Return no-op function for method calls
        return lambda *args, **kwargs: None


class StageProcessor(Protocol):
    """
    Protocol for stage processors.

    All stage adapters must implement this interface.
    """

    def process(
        self,
        config_path: Path,
        config_s3_path: str | None = None,
        dataset_name: str | None = None,
        job_id: str | None = None,
        status_updater: Any | None = None,
    ) -> dict[str, Any]:
        """
        Process stage with config file/directory.

        Args:
            config_path: Local path to config file or directory
            config_s3_path: S3 path to config (for result metadata)
            dataset_name: Optional dataset name override (from job metadata)
            job_id: Optional job ID for tracking
            status_updater: Optional status updater for progress tracking

        Returns:
            Result metadata dict

        Raises:
            RuntimeError: If processing fails
        """
        ...


class K8sStageAdapter:
    """
    Adapter for K8s MonthProcessor.

    Converts API job format to K8s engine format.
    """

    def __init__(self):
        """Initialize K8s stage adapter."""
        # Processor is created per-job for isolation
        pass

    def process(
        self,
        config_path: Path,
        config_s3_path: str | None = None,
        dataset_name: str | None = None,
        job_id: str | None = None,
        status_updater: Any | None = None,
    ) -> dict[str, Any]:
        """
        Run K8s allocation stage.

        Args:
            config_path: Path to k8s config file
            dataset_name: Optional dataset name override (not used for k8s)
            job_id: Optional job ID for tracking
            status_updater: Optional status updater for progress tracking

        Returns:
            Result metadata for DynamoDB

        Raises:
            RuntimeError: If processing fails
        """
        start_time = time.time()

        try:
            # Track: Loading configuration
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "k8s_loading_config")

            # Load K8s configuration
            k8s_config = K8sConfig.from_yaml(str(config_path))
            k8s_config.validate()

            # Track: Dataset name from config
            if status_updater and job_id:
                status_updater.update_dataset_info(
                    job_id, k8s_config.dataset_name, "K8s allocation"
                )
                status_updater.add_processing_stage(job_id, "k8s_processing_months")

            # Create month processor (uses existing implementation)
            # Import NoOpProgressTracker for headless execution
            from k8s.services.progress_tracker import NoOpProgressTracker

            # Use NoOpConsole and NoOpProgressTracker for headless execution
            processor = MonthProcessor(
                k8s_config,
                console=NoOpConsole(),
                debug=False,
                progress_tracker=NoOpProgressTracker()
            )

            # Process all months (or specified month/provider from config)
            processing_summary = processor.process_all_months(
                target_month=None, target_provider=None
            )

            # Track: Processing complete
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "k8s_months_completed")

            # Update or create dataset metadata file
            if status_updater and job_id:
                from src.services.dataset_identification import DatasetMetadataService
                from src.services.s3_data_service import S3DataService

                s3_service = S3DataService()
                metadata_service = DatasetMetadataService(s3_service)

                # Try to update existing metadata, or create if dataset is new
                try:
                    description = status_updater.get_job_description(job_id)
                    metadata_service.update_basic_metadata(
                        dataset_id=k8s_config.dataset_name,
                        description=description
                    )
                except:
                    # If update fails, create new metadata
                    description = status_updater.get_job_description(job_id)
                    metadata_service.create_initial_metadata(
                        dataset_id=k8s_config.dataset_name,
                        description=description,
                        reference_data_version="Unknown"
                    )
                status_updater.add_processing_stage(job_id, "k8s_metadata_file_updated")

            # Calculate duration
            duration = int(time.time() - start_time)

            # Extract result metadata with per-month statistics
            return self._extract_result_metadata(
                config_path, config_s3_path, duration, processing_summary, k8s_config.dataset_name
            )

        except Exception as e:
            raise RuntimeError(f"K8s stage processing failed: {e}") from e

    def _extract_result_metadata(
        self, config_path: Path, config_s3_path: str | None, duration: int, processing_summary: dict[str, int], dataset_name: str
    ) -> dict[str, Any]:
        """
        Extract result metadata with per-month line item counts.

        Args:
            config_path: Config file path (local temp)
            config_s3_path: S3 path to config
            duration: Processing duration in seconds
            processing_summary: Summary from month processor
            dataset_name: Dataset identifier

        Returns:
            Result metadata dict with per-month statistics
        """
        # Get per-month statistics from S3
        per_month_stats = get_per_month_statistics(dataset_name, "k8s")

        result = {
            "stage": "k8s",
            "config_path": config_s3_path or str(config_path),  # Prefer S3 path
            "processing_duration_seconds": duration,
            "processed_months": processing_summary.get("processed_months", 0),
            "failed_months": processing_summary.get("failed_months", 0),
            "status": "success"
            if processing_summary.get("failed_months", 0) == 0
            else "partial",
        }

        # Add per-month statistics
        if per_month_stats["per_month"]:
            result["per_month_statistics"] = per_month_stats["per_month"]
            result["peak_month"] = per_month_stats["peak_month"]
            result["total_records"] = per_month_stats["total_records"]

        return result


# Generate Stage Adapter (Phase 1)
class GenerateStageAdapter:
    """
    Adapter for StreamingGenerationEngine.

    Converts API job format to Generate engine format.
    """

    def __init__(self):
        """Initialize Generate stage adapter."""
        # Processor is created per-job for isolation
        pass

    def process(
        self,
        config_path: Path,
        config_s3_path: str | None = None,
        dataset_name: str | None = None,
        job_id: str | None = None,
        status_updater: Any | None = None,
    ) -> dict[str, Any]:
        """
        Run Generate stage.

        Args:
            config_path: Path to generate config folder
            dataset_name: Optional dataset name override (not used for generate)
            job_id: Optional job ID for tracking
            status_updater: Optional status updater for progress tracking

        Returns:
            Result metadata for DynamoDB

        Raises:
            RuntimeError: If processing fails
        """
        import time
        from src.generate.config.models import load_generation_configs, merge_generation_configs
        from src.generate.engine import StreamingGenerationEngine
        from src.services.s3_data_service import S3DataService
        from datetime import datetime

        start_time = time.time()

        try:
            # Track: Loading configuration
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "generate_loading_config")

            # Load and merge configurations
            configs = load_generation_configs(config_path)
            if len(configs) > 1:
                merged_config = merge_generation_configs(configs)
            else:
                merged_config = configs[0]

            # Track: Config loaded
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "generate_config_loaded")

            # Generate dataset name
            dataset_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

            # Track: Dataset name created
            if status_updater and job_id:
                status_updater.update_dataset_info(
                    job_id, dataset_name, "Generate billing data"
                )
                status_updater.add_processing_stage(job_id, "generate_dataset_name_created")

            # Initialize S3 service and engine
            s3_service = S3DataService()

            # Use NoOpConsole for headless worker execution
            console = NoOpConsole()

            # Track: Initializing engine
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "generate_initializing_engine")

            engine = StreamingGenerationEngine(
                s3_service=s3_service,
                console=console,
                org_name=dataset_name
            )

            # Track: Starting generation
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "generate_starting_generation")
                status_updater.update_heartbeat(job_id)

            # Generate dataset
            results = engine.generate_dataset(
                merged_config,
                dataset_name,
                verbose=False
            )

            # Track: Generation complete
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "generate_generation_completed")

            # Create dataset metadata file with description
            if status_updater and job_id:
                from src.services.dataset_identification import DatasetMetadataService

                metadata_service = DatasetMetadataService(s3_service)
                description = status_updater.get_job_description(job_id)
                metadata_service.create_initial_metadata(
                    dataset_id=dataset_name,
                    description=description,
                    reference_data_version="Unknown"  # TODO: Track reference data version
                )
                status_updater.add_processing_stage(job_id, "generate_metadata_file_created")

            # Calculate duration
            duration = int(time.time() - start_time)

            # Extract result metadata with per-month statistics
            return self._extract_result_metadata(
                config_path, config_s3_path, dataset_name, duration, results
            )

        except Exception as e:
            import traceback
            print(f"ERROR: Full traceback for generate stage failure:")
            print(traceback.format_exc())
            raise RuntimeError(f"Generate stage processing failed: {e}") from e

    def _extract_result_metadata(
        self,
        config_path: Path,
        config_s3_path: str | None,
        dataset_name: str,
        duration: int,
        results: dict[str, Any],
    ) -> dict[str, Any]:
        """
        Extract result metadata with per-month line item counts.

        Args:
            config_path: Config folder path (local temp)
            config_s3_path: S3 path to config
            dataset_name: Generated dataset name
            duration: Processing duration in seconds
            results: Results from engine.generate_dataset()

        Returns:
            Result metadata dict with per-month statistics
        """
        # Actual results structure from StreamingGenerationEngine.generate_dataset():
        # {
        #     "dataset_name": str,
        #     "generation": {...},
        #     "line_items": {
        #         "summary": {"total_line_items": int, ...},
        #         "providers": {...}
        #     },
        #     "failed_jobs": [...],
        #     "success": bool
        # }

        # Extract line item stats
        line_items = results.get("line_items", {})
        summary = line_items.get("summary", {})
        total_records = summary.get("total_line_items", 0)

        # Count successful vs failed jobs from generation results
        generation = results.get("generation", {})
        successful_jobs = generation.get("successful_jobs", 0)
        failed_jobs_count = generation.get("failed_jobs", 0)

        # Get per-month statistics from S3
        per_month_stats = get_per_month_statistics(dataset_name, "generate")

        result = {
            "stage": "generate",
            "config_path": config_s3_path or str(config_path),  # Prefer S3 path
            "dataset_name": dataset_name,
            "processing_duration_seconds": duration,
            "total_records_generated": total_records,
            "successful_jobs": successful_jobs,
            "failed_jobs": failed_jobs_count,
            "status": "success" if results.get("success", True) else "partial",
        }

        # Add per-month statistics
        if per_month_stats["per_month"]:
            result["per_month_statistics"] = per_month_stats["per_month"]
            result["peak_month"] = per_month_stats["peak_month"]
            result["total_records"] = per_month_stats["total_records"]

        return result


# Placeholder adapters for other stages (to be implemented in Phase 2)


class MetadataStageAdapter:
    """
    Adapter for MetadataProcessor.

    Converts API job format to Metadata engine format.
    """

    def __init__(self):
        """Initialize Metadata stage adapter."""
        pass

    def process(
        self,
        config_path: Path,
        config_s3_path: str | None = None,
        dataset_name: str | None = None,
        job_id: str | None = None,
        status_updater: Any | None = None,
    ) -> dict[str, Any]:
        """
        Run Metadata allocation stage using MetadataOrchestrator.

        Args:
            config_path: Path to metadata config file
            dataset_name: Optional dataset name override (not used for metadata)
            job_id: Optional job ID for tracking
            status_updater: Optional status updater for progress tracking

        Returns:
            Result metadata for DynamoDB

        Raises:
            RuntimeError: If processing fails
        """
        start_time = time.time()

        try:
            # Track: Loading configuration
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "metadata_loading_config")

            # Import metadata modules (actual CLI implementation)
            from metadata.config import MetadataConfig
            from metadata.services.metadata_orchestrator import MetadataOrchestrator

            # Load configuration
            config = MetadataConfig.from_yaml(str(config_path))

            # Track: Dataset name from config
            if status_updater and job_id:
                status_updater.update_dataset_info(
                    job_id, config.dataset_name, "Metadata allocation"
                )
                status_updater.add_processing_stage(job_id, "metadata_config_loaded")

            # Initialize orchestrator (handles all S3 operations internally)
            orchestrator = MetadataOrchestrator(config)

            # Track: Starting metadata allocation
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "metadata_processing_started")
                status_updater.update_heartbeat(job_id)

            # Run metadata allocation - orchestrator handles:
            # - Month discovery from S3
            # - Provider/month iteration
            # - MetadataProcessor execution
            # - S3 writes with success markers
            result = orchestrator.run_metadata_allocation(
                provider=None,  # Process all providers
                month=None,     # Process all months
                force=False     # Don't reprocess completed months
            )

            # Track: Processing complete
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "metadata_processing_completed")

            # Update dataset metadata file
            if status_updater and job_id:
                from src.services.dataset_identification import DatasetMetadataService
                from src.services.s3_data_service import S3DataService

                s3_service = S3DataService()
                metadata_service = DatasetMetadataService(s3_service)

                try:
                    description = status_updater.get_job_description(job_id)
                    metadata_service.update_basic_metadata(
                        dataset_id=config.dataset_name,
                        description=description
                    )
                except:
                    description = status_updater.get_job_description(job_id)
                    metadata_service.create_initial_metadata(
                        dataset_id=config.dataset_name,
                        description=description,
                        reference_data_version="Unknown"
                    )
                status_updater.add_processing_stage(job_id, "metadata_metadata_file_updated")

            # Calculate duration
            duration = int(time.time() - start_time)

            # Check for errors
            if result.get("status") != "Completed":
                raise RuntimeError(f"Metadata processing did not complete: {result.get('status')}")

            stats = result.get("processing_stats", {})
            if stats.get("processing_errors", 0) > 0:
                raise RuntimeError(
                    f"Metadata processing had errors: {stats['processing_errors']} failures"
                )

            # Extract result metadata
            return self._extract_result_metadata(config_path, duration, result)

        except Exception as e:
            raise RuntimeError(f"Metadata stage processing failed: {e}") from e

    def _extract_result_metadata(
        self, config_path: Path, duration: int, result: dict[str, Any]
    ) -> dict[str, Any]:
        """
        Extract result metadata from orchestrator result.

        Args:
            config_path: Config file path
            duration: Processing duration in seconds
            result: Result from orchestrator.run_metadata_allocation()

        Returns:
            Result metadata dict
        """
        stats = result.get("processing_stats", {})
        coverage = result.get("coverage", {})

        return {
            "stage": "metadata",
            "config_path": str(config_path),
            "processing_duration_seconds": duration,
            "providers_processed": coverage.get("providers_processed", []),
            "months_processed": coverage.get("months_processed", []),
            "total_records": stats.get("total_rows_processed", 0),
            "processed_successfully": stats.get("processed_successfully", 0),
            "status": "success",
        }


class UsageStageAdapter:
    """
    Adapter for Usage stage using process_usage_allocation.

    Converts API job format to Usage CLI format and uses the actual
    CLI processor for processing. Follows K8sStageAdapter pattern.
    """

    def __init__(self):
        """Initialize Usage stage adapter."""
        pass

    def process(
        self,
        config_path: Path,
        config_s3_path: str | None = None,
        dataset_name: str | None = None,
        job_id: str | None = None,
        status_updater: Any | None = None,
    ) -> dict[str, Any]:
        """
        Run Usage allocation stage using process_usage_allocation.

        Args:
            config_path: Path to usage config file
            dataset_name: Optional dataset name override (not used for usage)
            job_id: Optional job ID for tracking
            status_updater: Optional status updater for progress tracking

        Returns:
            Result metadata for DynamoDB

        Raises:
            RuntimeError: If processing fails
        """
        start_time = time.time()

        try:
            # Track: Loading configuration
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "usage_loading_config")

            # Import usage modules (actual CLI implementation)
            from usage.config.models import load_usage_config
            from usage.core.processor import process_usage_allocation

            # Load configuration
            config = load_usage_config(str(config_path))

            # Track: Dataset name from config
            if status_updater and job_id:
                status_updater.update_dataset_info(
                    job_id, config.dataset_name, "Usage allocation"
                )
                status_updater.add_processing_stage(job_id, "usage_config_loaded")

            # Ensure reference data is available (download from S3 if needed)
            from api.worker.reference_manager import ensure_reference_data
            s3_bucket = os.environ.get("CZBILLGEN_S3_BUCKET")
            if not s3_bucket:
                raise RuntimeError("CZBILLGEN_S3_BUCKET environment variable is required")

            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "usage_checking_reference_data")

            if not ensure_reference_data(s3_bucket):
                raise RuntimeError("Failed to download reference data from S3")

            # Track: Starting usage allocation
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "usage_processing_started")
                status_updater.update_heartbeat(job_id)

            # Process usage allocation - processor handles:
            # - Month discovery from S3 (metadata stage)
            # - Reference data loading
            # - Usage calculation
            # - S3 writes with success markers
            result = process_usage_allocation(
                config=config,
                specific_month=None,  # Process all months
                force=False           # Don't reprocess completed months
            )

            # Track: Processing complete
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "usage_processing_completed")

            # Update dataset metadata file
            if status_updater and job_id:
                from src.services.dataset_identification import DatasetMetadataService
                from src.services.s3_data_service import S3DataService

                s3_service = S3DataService()
                metadata_service = DatasetMetadataService(s3_service)

                try:
                    description = status_updater.get_job_description(job_id)
                    metadata_service.update_basic_metadata(
                        dataset_id=config.dataset_name,
                        description=description
                    )
                except:
                    description = status_updater.get_job_description(job_id)
                    metadata_service.create_initial_metadata(
                        dataset_id=config.dataset_name,
                        description=description,
                        reference_data_version="Unknown"
                    )
                status_updater.add_processing_stage(job_id, "usage_metadata_file_updated")

            # Calculate duration
            duration = int(time.time() - start_time)

            # Check for errors
            if not result.success:
                raise RuntimeError(f"Usage processing failed: {result.error}")

            # Extract result metadata
            return self._extract_result_metadata(config_path, duration, result)

        except Exception as e:
            raise RuntimeError(f"Usage stage processing failed: {e}") from e

    def _extract_result_metadata(
        self, config_path: Path, duration: int, result: Any
    ) -> dict[str, Any]:
        """
        Extract result metadata from ProcessingResult.

        Args:
            config_path: Config file path
            duration: Processing duration in seconds
            result: ProcessingResult from process_usage_allocation()

        Returns:
            Result metadata dict
        """
        return {
            "stage": "usage",
            "config_path": str(config_path),
            "processing_duration_seconds": duration,
            "processed_months": result.processed_months,
            "total_records": result.total_records,
            "usage_records": result.usage_records,
            "usage_coverage": result.usage_coverage,
            "status": "success",
        }


class DiscountStageAdapter:
    """
    Adapter for Discount stage using process_discount_allocation.

    Converts API job format to Discount CLI format and uses the actual
    CLI processor for processing. Follows K8sStageAdapter pattern.
    """

    def __init__(self):
        """Initialize Discount stage adapter."""
        pass

    def process(
        self,
        config_path: Path,
        config_s3_path: str | None = None,
        dataset_name: str | None = None,
        job_id: str | None = None,
        status_updater: Any | None = None,
    ) -> dict[str, Any]:
        """
        Run Discount allocation stage using process_discount_allocation.

        Args:
            config_path: Path to discount config file
            dataset_name: Optional dataset name override (not used for discount)
            job_id: Optional job ID for tracking
            status_updater: Optional status updater for progress tracking

        Returns:
            Result metadata for DynamoDB

        Raises:
            RuntimeError: If processing fails
        """
        start_time = time.time()

        try:
            # Track: Loading configuration
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "discount_loading_config")

            # Import discount modules (actual CLI implementation)
            from discount.config.models import load_discount_config
            from discount.core.processor import process_discount_allocation

            # Load configuration
            config = load_discount_config(str(config_path))

            # Track: Dataset name from config
            if status_updater and job_id:
                status_updater.update_dataset_info(
                    job_id, config.dataset_name, "Discount allocation"
                )
                status_updater.add_processing_stage(job_id, "discount_config_loaded")

            # Track: Starting discount allocation
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "discount_processing_started")
                status_updater.update_heartbeat(job_id)

            # Process discount allocation - processor handles:
            # - Dual-source processing (usage stage first, then metadata gaps)
            # - Processor factory with selective activation
            # - Schema harmonization
            # - S3 writes with success markers
            result = process_discount_allocation(
                config=config,
                specific_month=None,  # Process all months
                force=False           # Don't reprocess completed months
            )

            # Track: Processing complete
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "discount_processing_completed")

            # Update dataset metadata file
            if status_updater and job_id:
                from src.services.dataset_identification import DatasetMetadataService
                from src.services.s3_data_service import S3DataService

                s3_service = S3DataService()
                metadata_service = DatasetMetadataService(s3_service)

                try:
                    description = status_updater.get_job_description(job_id)
                    metadata_service.update_basic_metadata(
                        dataset_id=config.dataset_name,
                        description=description
                    )
                except:
                    description = status_updater.get_job_description(job_id)
                    metadata_service.create_initial_metadata(
                        dataset_id=config.dataset_name,
                        description=description,
                        reference_data_version="Unknown"
                    )
                status_updater.add_processing_stage(job_id, "discount_metadata_file_updated")

            # Calculate duration
            duration = int(time.time() - start_time)

            # Check for errors
            if not result.success:
                raise RuntimeError(f"Discount processing failed: {result.error}")

            # Extract result metadata
            return self._extract_result_metadata(config_path, duration, result)

        except Exception as e:
            raise RuntimeError(f"Discount stage processing failed: {e}") from e

    def _extract_result_metadata(
        self, config_path: Path, duration: int, result: Any
    ) -> dict[str, Any]:
        """
        Extract result metadata from ProcessingResult.

        Args:
            config_path: Config file path
            duration: Processing duration in seconds
            result: ProcessingResult from process_discount_allocation()

        Returns:
            Result metadata dict
        """
        return {
            "stage": "discount",
            "config_path": str(config_path),
            "processing_duration_seconds": duration,
            "processed_months": result.processed_months,
            "usage_months": result.usage_months,
            "metadata_gaps": result.metadata_gaps,
            "status": "success",
        }


class UploadGeneratedAdapter:
    """
    Adapter for multi-org dataset uploads.

    Converts API job format to Upload workflow format.
    Follows same pattern as processing stage adapters.
    """

    def __init__(self):
        """Initialize Upload adapter."""
        pass

    def process(
        self,
        config_path: Path,
        config_s3_path: str | None = None,
        dataset_name: str | None = None,
        job_id: str | None = None,
        status_updater: Any | None = None,
        recreate_connections: bool = False,
    ) -> dict[str, Any]:
        """
        Run upload workflow for multiple CloudZero organizations.

        Args:
            config_path: Path to upload config YAML (downloaded from S3)
            dataset_name: Dataset to upload (overrides config if provided)
            job_id: Job ID for tracking
            status_updater: StatusUpdater for progress tracking
            recreate_connections: Delete and recreate CloudZero connections

        Returns:
            Result metadata for DynamoDB

        Raises:
            RuntimeError: If upload fails
        """
        start_time = time.time()

        try:
            config = self._load_and_validate_config(
                config_path, dataset_name, job_id, status_updater
            )
            result = self._execute_upload(config, job_id, status_updater, recreate_connections)
            self._track_upload_results(result, job_id, status_updater)

            duration = int(time.time() - start_time)

            if result.status == 'failed':
                error = result.error_message or "Unknown upload failure"
                raise RuntimeError(f"Upload failed: {error}")

            return self._extract_result_metadata(
                config_path, duration, config, result
            )

        except Exception as e:
            raise RuntimeError(f"Upload stage processing failed: {e}") from e

    def _load_and_validate_config(
        self,
        config_path: Path,
        config_s3_path: str | None = None,
        dataset_name: str | None = None,
        job_id: str | None = None,
        status_updater: Any | None = None,
    ) -> Any:
        """Load and validate upload config."""
        if status_updater and job_id:
            status_updater.add_processing_stage(job_id, "upload_loading_config")

        from src.upload.models import UploadConfig

        upload_config = UploadConfig.from_yaml(config_path)

        if dataset_name:
            upload_config.dataset_name = dataset_name

        if status_updater and job_id:
            status_updater.update_dataset_info(
                job_id,
                upload_config.dataset_name,
                f"Upload to {len(upload_config.organizations)} organizations"
            )
            status_updater.add_processing_stage(job_id, "upload_config_loaded")

        return upload_config

    def _execute_upload(
        self,
        upload_config: Any,
        job_id: str | None,
        status_updater: Any | None,
        recreate_connections: bool = False,
    ) -> Any:
        """Execute upload workflow via coordinator."""
        if status_updater and job_id:
            stage_name = "upload_recreate_starting" if recreate_connections else "upload_starting"
            status_updater.add_processing_stage(job_id, stage_name)
            status_updater.update_heartbeat(job_id)

        from src.upload.services import UploadServiceFactory

        coordinator = UploadServiceFactory.create_coordinator(upload_config)

        if recreate_connections:
            return coordinator.recreate_connections_and_upload(upload_config)
        else:
            return coordinator.upload_dataset(upload_config, dry_run=False)

    def _track_upload_results(
        self,
        result: Any,
        job_id: str | None,
        status_updater: Any | None,
    ) -> None:
        """Track per-organization upload results with progress."""
        if not status_updater or not job_id:
            return

        total_orgs = len(result.results)
        current_org = 0

        for org_name, org_result in result.results.items():
            current_org += 1
            org_status, files, providers = self._extract_org_result(org_result)
            stage_name = f"upload_org_{org_name}_{org_status}"

            # Calculate progress percentage
            progress = {
                "current": current_org,
                "total": total_orgs,
                "percent": (current_org / total_orgs * 100) if total_orgs > 0 else 0
            }

            status_updater.add_processing_stage(
                job_id,
                stage_name,
                progress=progress,
                status="completed"
            )

        complete_stage = f"upload_complete_{result.organizations_successful}_of_{result.organizations_processed}"
        status_updater.add_processing_stage(job_id, complete_stage)

    def _extract_org_result(self, org_result: Any) -> tuple[str, int, list]:
        """Extract status, files, and providers from org result."""
        if hasattr(org_result, 'status'):
            org_status = org_result.status
            files = getattr(org_result, 'files_uploaded', 0)
            providers = list(getattr(org_result, 'providers', []))
        else:
            org_status = org_result.get('status', 'unknown')
            files = org_result.get('files_uploaded', 0)
            providers = org_result.get('providers', [])

        return org_status, files, providers

    def _extract_result_metadata(
        self, config_path: Path, duration: int, config: Any, result: Any
    ) -> dict[str, Any]:
        """
        Extract result metadata from WorkflowResult.

        Args:
            config_path: Config file path
            duration: Processing duration in seconds
            config: UploadConfig
            result: WorkflowResult from coordinator

        Returns:
            Result metadata dict
        """
        # Calculate total files uploaded
        total_files = 0
        for org_result in result.results.values():
            if hasattr(org_result, 'files_uploaded'):
                total_files += org_result.files_uploaded
            elif isinstance(org_result, dict):
                total_files += org_result.get('files_uploaded', 0)

        return {
            "stage": "upload_generated",
            "config_path": str(config_path),
            "processing_duration_seconds": duration,
            "dataset": result.dataset,
            "organizations_successful": result.organizations_successful,
            "organizations_processed": result.organizations_processed,
            "files_uploaded": total_files,
            "status": result.status,
        }


class TimeShiftAdapter:
    """
    Adapter for time-shift jobs in ECS worker.

    Executes TimeShiftOrchestrator with DynamoDB status tracking
    and breadcrumb trail for progress monitoring.
    """

    def __init__(self):
        """Initialize time-shift adapter."""
        pass

    def process(
        self,
        config_path: Path,
        config_s3_path: str | None = None,
        dataset_name: str | None = None,
        job_id: str | None = None,
        status_updater: Any | None = None,
    ) -> dict[str, Any]:
        """
        Execute time-shift job.

        Transforms existing billing drops to new time periods via:
        1. Reading source CSV.gz files from S3
        2. Transforming timestamps in memory (Polars)
        3. Writing to target month folders in S3

        Args:
            config_path: Path to time-shift config YAML
            dataset_name: Not used for time-shift
            job_id: Job identifier for tracking
            status_updater: Status updater for DynamoDB tracking

        Returns:
            Result dictionary with statistics

        Raises:
            RuntimeError: If any drops fail
        """
        start_time = time.time()

        try:
            # Track: Loading configuration
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "timeshift_loading_config")

            # Import time-shift services
            from src.upload.services import (
                TimeShiftConfigLoader,
                TargetMonthParser,
                SourceMonthSelector,
                DayCountCalculator,
                CsvDateTransformer,
                ManifestManager,
                BillingPeriodHelper,
                DropCopier,
                TimeShiftProviderCopier,
                TimeShiftOrchestrator
            )
            from src.services.environment_service import EnvironmentService

            # Load environment and config
            env_service = EnvironmentService()
            s3_bucket = env_service.require(
                "CZBILLGEN_S3_BUCKET",
                "CZBILLGEN_S3_BUCKET not configured"
            )
            s3_root_folder = "piedpiper-custom-billing"

            config_loader = TimeShiftConfigLoader()
            ts_config = config_loader.load_config(str(config_path))

            # Track: Configuration loaded
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "timeshift_config_loaded")

            # Initialize services with dependency injection
            source_selector = SourceMonthSelector(
                s3_bucket=s3_bucket,
                s3_root_folder=s3_root_folder
            )
            day_calculator = DayCountCalculator()
            csv_transformer = CsvDateTransformer(day_calculator)
            manifest_manager = ManifestManager(s3_bucket)
            billing_helper = BillingPeriodHelper()
            drop_copier = DropCopier(
                csv_transformer=csv_transformer,
                manifest_manager=manifest_manager,
                billing_helper=billing_helper,
                s3_bucket=s3_bucket,
                s3_root_folder=s3_root_folder
            )
            provider_copier = TimeShiftProviderCopier(
                manifest_manager=manifest_manager,
                billing_helper=billing_helper,
                s3_bucket=s3_bucket,
                s3_root_folder=s3_root_folder
            )
            target_parser = TargetMonthParser()
            orchestrator = TimeShiftOrchestrator(
                config_loader=config_loader,
                target_parser=target_parser,
                source_selector=source_selector,
                drop_copier=drop_copier,
                provider_copier=provider_copier
            )

            # Track: Starting time-shift
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "timeshift_processing_started")
                status_updater.update_heartbeat(job_id)

            # Execute time-shift
            result = orchestrator.execute(str(config_path))

            # Track: Processing complete
            if status_updater and job_id:
                status_updater.add_processing_stage(job_id, "timeshift_processing_completed")

            # Calculate duration
            duration = int(time.time() - start_time)

            # Check for failures
            failed_results = [dr for dr in result.drop_results if not dr.success]
            if failed_results:
                error_details = "; ".join([
                    f"{dr.org}/{dr.provider}/{dr.target_month}: {dr.error_message}"
                    for dr in failed_results[:5]
                ])
                raise RuntimeError(
                    f"Time-shift had {len(failed_results)} failures: {error_details}"
                )

            # Extract result metadata
            return self._extract_result_metadata(config_path, duration, result)

        except Exception as e:
            raise RuntimeError(f"Time-shift processing failed: {e}") from e

    def _extract_result_metadata(
        self, config_path: Path, duration: int, result: Any
    ) -> dict[str, Any]:
        """
        Extract result metadata from orchestrator result.

        Args:
            config_path: Config file path
            duration: Processing duration in seconds
            result: TimeShiftResult from orchestrator.execute()

        Returns:
            Result metadata dict
        """
        return {
            "stage": "time_shift",
            "config_path": str(config_path),
            "processing_duration_seconds": duration,
            "drops_successful": result.drops_processed,
            "drops_total": result.drops_processed + result.drops_failed,
            "total_records": result.total_records,
            "status": "success",
        }
