"""Month-by-month processing coordinator for usage allocation."""

import gc
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
import polars as pl

from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn

from ..config.models import UsageConfig
from src.usage.services.usage_data_service import UsageDataService
from .calculator import calculate_usage_amounts
from .reference_loader import load_reference_data_for_provider, clear_reference_cache
from .validation import validate_before_write
from src.log_context.core import get_logger

console = Console()
logger = get_logger(__name__)


@dataclass
class ProcessingResult:
    """Result of usage allocation processing."""
    success: bool
    error: Optional[str] = None
    processed_months: int = 0
    usage_records: int = 0
    total_records: int = 0
    usage_coverage: float = 0.0
    message: Optional[str] = None


def process_usage_allocation(
    config: UsageConfig,
    specific_month: Optional[str] = None,
    force: bool = False
) -> ProcessingResult:
    """
    Process usage allocation with fail-fast behavior.
    
    Args:
        config: Usage configuration
        specific_month: Optional specific month to process (YYYY-MM)
        force: Force reprocessing of completed months
        
    Returns:
        ProcessingResult with success status and metrics
    """
    try:
        console.print(f"🚀 Starting usage allocation for dataset: {config.dataset_name}")
        
        if specific_month:
            console.print(f"📅 Processing specific month: {specific_month}")
        
        if force:
            console.print("🔄 Force mode enabled: Will reprocess completed months")
        
        # Process all months with fail-fast
        return process_all_months_fail_fast(config, specific_month, force)
        
    except Exception as e:
        logger.error(f"Usage allocation failed: {e}")
        console.print(f"❌ Processing failed: {e}")
        return ProcessingResult(
            success=False,
            error=str(e)
        )


def process_all_months_fail_fast(
    config: UsageConfig,
    specific_month: Optional[str] = None,
    force: bool = False
) -> ProcessingResult:
    """
    Process all months with fail-fast behavior.
    
    If ANY month fails for ANY reason (validation, processing, S3), 
    immediately stop and report failure. Clean up any partial data.
    
    Args:
        config: Usage configuration
        specific_month: Optional specific month to process
        force: Force reprocessing of completed months
        
    Returns:
        ProcessingResult with success status and metrics
    """
    
    data_service = UsageDataService()

    # Discover months to process (respecting --force)
    if specific_month:
        months_to_process = {provider: [specific_month] for provider in discover_providers(config.dataset_name)}
        # Filter out providers that don't have the specific month
        months_to_process = {p: months for p, months in months_to_process.items() if months}
        source_stage = "unknown"  # Specific month mode
    else:
        months_to_process, source_stage = discover_months_to_process(config.dataset_name, force, data_service)

    # Distinguish between no input data vs already processed
    if not months_to_process:
        if source_stage is None:
            # ERROR: No input data found
            return ProcessingResult(
                success=False,
                error="No input data found in any pipeline stage"
            )
        else:
            # SUCCESS: Everything already processed
            return ProcessingResult(
                success=True,
                message="All months already processed"
            )

    # Handle --force: Delete existing data before processing
    if force:
        delete_existing_usage_data(config.dataset_name, months_to_process, data_service)

    total_months = sum(len(months) for months in months_to_process.values())
    console.print(f"🚀 Processing {total_months} months from {source_stage} stage")
    
    # Process months with fail-fast
    reference_cache = {}
    processed_successfully = []
    skipped_providers = []
    skipped_months = []  # Track individual skipped months
    total_usage_records = 0
    total_records = 0
    
    try:
        # Use rich progress tracking
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TaskProgressColumn(),
            console=console,
            transient=False
        ) as progress:
            
            # Create main progress task
            main_task = progress.add_task("Processing months...", total=total_months)
            
            processed_count = 0
            
            for provider, months in months_to_process.items():
                # Load reference data once per provider
                if provider not in reference_cache:
                    progress.update(main_task, description=f"Loading reference data for {provider}...")
                    try:
                        reference_data = load_reference_data_for_provider(provider)
                        
                        # Skip provider if no pricing data available
                        if reference_data.is_empty():
                            console.print(f"⏭️  Skipping {provider}: No pricing data available")
                            logger.info(f"Skipping provider {provider} - no pricing data available")
                            skipped_providers.append(provider)
                            # Update progress for all months of this provider
                            progress.update(main_task, advance=len(months))
                            continue
                        
                        reference_cache[provider] = reference_data
                        
                    except Exception as e:
                        # FAIL FAST: Reference data loading failed
                        console.print(f"❌ FATAL: Failed to load reference data for {provider}: {e}")
                        cleanup_partial_processing(config.dataset_name, processed_successfully, data_service)
                        return ProcessingResult(success=False, error=f"Reference data failed: {e}")
                
                # Skip if provider has no pricing data (already skipped above, but handle cached case)
                if provider in reference_cache and reference_cache[provider].is_empty():
                    console.print(f"⏭️  Skipping {provider}: No pricing data available (cached)")
                    if provider not in skipped_providers:
                        skipped_providers.append(provider)
                    progress.update(main_task, advance=len(months))
                    continue
                
                for month in months:
                    progress.update(main_task, description=f"Processing {provider} {month}...")
                    
                    try:
                        # Process month with validation
                        result = process_month_with_validation(
                            config.dataset_name, provider, month, reference_cache[provider], data_service, source_stage
                        )
                        
                        if not result.success:
                            # FAIL FAST: Month processing failed
                            console.print(f"❌ FATAL: {provider} {month} failed: {result.error}")
                            # Only cleanup for critical failures, not validation failures
                            if "Critical validation failed" in result.error:
                                cleanup_partial_processing(config.dataset_name, processed_successfully, data_service)
                            return ProcessingResult(success=False, error=result.error)
                        
                        # Check if month was skipped (zero usage records but success=True)
                        if result.usage_records == 0 and result.total_records > 0:
                            # Month was skipped due to no matches
                            skipped_months.append((provider, month))
                            total_records += result.total_records
                            processed_count += 1
                            
                            # Update progress
                            progress.update(main_task, advance=1)
                            
                            # Month skip message
                            console.print(f"⏭️  {provider} {month}: No usage matches found - skipped")
                        else:
                            # Success with usage records
                            processed_successfully.append((provider, month))
                            total_usage_records += result.usage_records
                            total_records += result.total_records
                            processed_count += 1
                            
                            # Update progress
                            progress.update(main_task, advance=1)
                            
                            # Month completion message with metrics
                            usage_rate = f"{result.usage_coverage:.1%}"
                            console.print(f"✅ {provider} {month}: {result.usage_records:,}/{result.total_records:,} usage records ({usage_rate})")
                        
                        # Memory management: Force garbage collection between months
                        del result
                        gc.collect()
                        
                    except Exception as e:
                        # FAIL FAST: Unexpected error
                        console.print(f"❌ FATAL: Unexpected error processing {provider} {month}: {e}")
                        # Only cleanup for unexpected system errors, not business logic failures
                        cleanup_partial_processing(config.dataset_name, processed_successfully, data_service)
                        return ProcessingResult(success=False, error=f"Unexpected error: {e}")
        
        # ALL months succeeded - write SUCCESS marker
        data_service.write_success_marker(config.dataset_name, {
            "status": "complete",
            "validated": True,
            "processed_months": len(processed_successfully),
            "force_used": force
        })
        
        overall_coverage = (total_usage_records / total_records * 100) if total_records > 0 else 0
        
        # Build completion message
        completion_parts = [f"{len(processed_successfully)} months processed"]
        if overall_coverage > 0:
            completion_parts.append(f"{overall_coverage:.1f}% coverage")
        if skipped_providers:
            completion_parts.append(f"{len(skipped_providers)} providers skipped (no pricing data)")
        if skipped_months:
            completion_parts.append(f"{len(skipped_months)} months skipped (no matches)")
        
        completion_message = "🎉 Usage allocation complete: " + ", ".join(completion_parts)
        console.print(completion_message)
        
        if skipped_providers:
            console.print(f"ℹ️  Skipped providers: {', '.join(skipped_providers)}")
        if skipped_months:
            console.print(f"ℹ️  Skipped months (no matches): {len(skipped_months)} total")
        
        return ProcessingResult(
            success=True, 
            processed_months=len(processed_successfully),
            usage_records=total_usage_records,
            total_records=total_records,
            usage_coverage=overall_coverage / 100,
            message=f"Processed {len(processed_successfully)} months, skipped {len(skipped_providers)} providers, {len(skipped_months)} months"
        )
        
    except KeyboardInterrupt:
        console.print("\n⚠️  Processing interrupted by user")
        cleanup_partial_processing(config.dataset_name, processed_successfully, data_service)
        return ProcessingResult(success=False, error="Interrupted by user")
    
    finally:
        # Always clean up reference cache
        clear_reference_cache()


def process_month_with_validation(
    dataset_name: str,
    provider: str,
    month: str,
    reference_data: pl.DataFrame,
    data_service: UsageDataService,
    source_stage: str = "metadata"
) -> ProcessingResult:
    """Process a single month with usage allocation and validation.

    Args:
        dataset_name: Dataset identifier
        provider: Cloud provider
        month: Month in YYYY-MM format
        reference_data: Pre-loaded reference data for provider
        data_service: S3 data service
        source_stage: Pipeline stage to load from (default: "metadata")

    Returns:
        ProcessingResult with success status and metrics
    """

    try:
        # 1. Download input data from specified stage
        original_data = data_service.download_month_data(dataset_name, provider, month, source_stage)
        
        # 2. Apply usage calculation
        enhanced_data = calculate_usage_amounts(original_data, reference_data)
        
        # 3. Pre-write validation (CRITICAL GATE)
        validation_result = validate_before_write(original_data, enhanced_data, provider, month)
        
        if validation_result.should_skip:
            logger.info(
                "Skipping month - no usage matches found",
                provider=provider,
                month=month
            )
            # Return success with zero records to indicate skip
            return ProcessingResult(
                success=True,
                usage_records=0,
                total_records=len(original_data),
                usage_coverage=0.0
            )
        
        if not validation_result.passed:
            logger.error(
                "Pre-write validation failed - data NOT written to S3",
                provider=provider,
                month=month,
                error=validation_result.error
            )
            return ProcessingResult(success=False, error=validation_result.error)
        
        # 4. Write to S3 only if validation passes
        upload_success = data_service.upload_month_data(dataset_name, provider, month, enhanced_data)
        
        if not upload_success:
            return ProcessingResult(success=False, error="S3 upload failed")
        
        # 5. Calculate metrics before cleanup
        usage_count = len(enhanced_data.filter(pl.col("usage_amount") != "0.0"))
        total_count = len(enhanced_data)
        coverage = usage_count / total_count if total_count > 0 else 0
        
        # 6. Memory cleanup - explicit deletion of large DataFrames
        del original_data
        del enhanced_data
        gc.collect()
        
        return ProcessingResult(
            success=True,
            usage_records=usage_count,
            total_records=total_count,
            usage_coverage=coverage
        )
        
    except Exception as e:
        logger.error(f"Month processing failed: {e}", provider=provider, month=month)
        return ProcessingResult(success=False, error=str(e))


def filter_unprocessed_months(
    dataset_name: str,
    all_months: Dict[str, List[str]],
    data_service: UsageDataService
) -> Dict[str, List[str]]:
    """
    Filter out months that are already processed.

    Single Responsibility: Month filtering only.

    Args:
        dataset_name: Dataset identifier
        all_months: All available months by provider
        data_service: S3 data service

    Returns:
        Dict[provider, List[months]] that need processing
    """
    months_to_process = {}
    skipped_count = 0

    for provider, months in all_months.items():
        unprocessed = []

        for month in months:
            if is_month_already_processed(dataset_name, provider, month, data_service):
                console.print(f"⏭️  Skipping {provider} {month}: Already processed")
                skipped_count += 1
            else:
                unprocessed.append(month)

        if unprocessed:
            months_to_process[provider] = unprocessed

    # Summary message
    total_to_process = sum(len(m) for m in months_to_process.values())
    if skipped_count > 0:
        console.print(
            f"📊 Processing {total_to_process} months, "
            f"skipped {skipped_count} already completed"
        )

    return months_to_process


def discover_months_to_process(
    dataset_name: str,
    force: bool = False,
    data_service: UsageDataService = None
) -> Tuple[Dict[str, List[str]], Optional[str]]:
    """
    Discover months with clear messaging about source and state.

    Args:
        dataset_name: Dataset identifier
        force: Force reprocessing of completed months
        data_service: S3 data service

    Returns:
        Tuple of (months_to_process, source_stage)
        - months_to_process: Dict[provider, List[months]]
        - source_stage: Stage where data was found, or None if no input data
    """

    if data_service is None:
        data_service = UsageDataService()

    # Get months with fallback
    all_available_months, source_stage = data_service.list_available_months_with_fallback(
        dataset_name
    )

    # Handle no input data case
    if not all_available_months:
        console.print("❌ No input data found in any pipeline stage")
        console.print("💡 Expected stages: metadata → k8s → generate")
        console.print("💡 Run 'czbg generate' to create billing data")
        return {}, None

    # Show which stage we're using
    console.print(f"📂 Using input data from: {source_stage} stage")
    if source_stage != "metadata":
        console.print(f"ℹ️  Metadata stage not yet run - using {source_stage} data directly")

    # Force mode - return everything
    if force:
        console.print("🔄 --force flag: Will reprocess all available months")
        return all_available_months, source_stage

    # Filter out already processed months
    months_to_process = filter_unprocessed_months(
        dataset_name, all_available_months, data_service
    )

    # No months to process - distinguish why
    if not months_to_process:
        console.print("✅ All available months already processed in usage stage")
        console.print("💡 Use --force to reprocess completed months")
        return {}, source_stage  # Note: source_stage is NOT None here

    return months_to_process, source_stage


def discover_providers(dataset_name: str) -> List[str]:
    """Discover available providers for specific month processing.
    
    Args:
        dataset_name: Dataset identifier
        
    Returns:
        List of available providers
    """
    try:
        data_service = UsageDataService()
        months_by_provider = data_service.list_available_months(dataset_name)
        return list(months_by_provider.keys())
    except Exception as e:
        logger.warning(f"Could not discover providers: {e}")
        return ["AWS", "GCP", "Azure"]  # Default fallback


def is_month_already_processed(
    dataset_name: str, 
    provider: str, 
    month: str,
    data_service: UsageDataService = None
) -> bool:
    """
    Check if month is already processed by checking S3 output exists.
    
    Args:
        dataset_name: Dataset identifier
        provider: Cloud provider
        month: Month in YYYY-MM format
        data_service: S3 data service
        
    Returns:
        True if month is already processed
    """
    if data_service is None:
        data_service = UsageDataService()
    
    try:
        s3_path = f"czbillgen/datasets/{dataset_name}/usage/{provider}/{month}.parquet"
        return data_service.file_exists(s3_path)
    except Exception:
        return False  # If we can't check, assume not processed


def cleanup_partial_processing(
    dataset_name: str, 
    processed_successfully: List[Tuple[str, str]],
    data_service: UsageDataService
) -> None:
    """Clean up partially processed data when fail-fast is triggered.
    
    Args:
        dataset_name: Dataset identifier
        processed_successfully: List of (provider, month) tuples that were processed
        data_service: S3 data service
    """
    
    if not processed_successfully:
        return
    
    console.print(f"🧹 Cleaning up {len(processed_successfully)} partially processed months...")
    
    for provider, month in processed_successfully:
        s3_path = f"czbillgen/datasets/{dataset_name}/usage/{provider}/{month}.parquet"
        try:
            data_service.delete_file(s3_path)
        except Exception as e:
            console.print(f"⚠️  Warning: Could not clean up {s3_path}: {e}")
    
    console.print("🧹 Cleanup complete - no partial data left in S3")


def delete_existing_usage_data(
    dataset_name: str, 
    months_to_process: Dict[str, List[str]],
    data_service: UsageDataService
) -> None:
    """Delete existing usage data when --force is used.
    
    Args:
        dataset_name: Dataset identifier
        months_to_process: Dict of provider -> months to delete
        data_service: S3 data service
    """
    
    console.print("🗑️  --force: Deleting existing usage data...")
    
    deleted_count = 0
    for provider, months in months_to_process.items():
        for month in months:
            s3_path = f"czbillgen/datasets/{dataset_name}/usage/{provider}/{month}.parquet"
            try:
                if data_service.file_exists(s3_path):
                    data_service.delete_file(s3_path)
                    deleted_count += 1
            except Exception as e:
                console.print(f"⚠️  Warning: Could not delete {s3_path}: {e}")
    
    # Also delete SUCCESS marker if it exists
    success_path = f"czbillgen/datasets/{dataset_name}/usage/_SUCCESS"
    try:
        if data_service.file_exists(success_path):
            data_service.delete_file(success_path)
            console.print("🗑️  Deleted usage stage SUCCESS marker")
    except Exception as e:
        console.print(f"⚠️  Warning: Could not delete SUCCESS marker: {e}")
    
    if deleted_count > 0:
        console.print(f"🗑️  Deleted {deleted_count} existing month files")