"""Main discount processing coordinator with dual-source support."""

from dataclasses import dataclass
from typing import Dict, List, Optional
from rich.console import Console
import polars as pl

from ..config.models import DiscountConfig
from ..services.discount_data_service import DiscountDataService
from .gap_analyzer import GapAnalyzer, ProcessingGap
from .schema_harmonizer import SchemaHarmonizer
from .factory import apply_discount_processors, ProcessorFactory
from ...log_context.core import get_logger

console = Console()


@dataclass
class ProcessingResult:
    """Result of discount processing operation."""
    success: bool
    processed_months: int = 0
    usage_months: int = 0
    metadata_gaps: int = 0
    error: Optional[str] = None


def process_discount_allocation(
    config: DiscountConfig,
    specific_month: Optional[str] = None,
    force: bool = False
) -> ProcessingResult:
    """
    Main processing coordinator implementing usage -> metadata flow.
    
    Args:
        config: Discount configuration
        specific_month: Optional specific month to process
        force: Force reprocessing of completed months
        
    Returns:
        ProcessingResult with success status and metrics
    """
    logger = get_logger(__name__)
    data_service = DiscountDataService()
    gap_analyzer = GapAnalyzer()
    
    try:
        logger.info(
            "Starting discount allocation with dual-source processing",
            dataset_name=config.dataset_name,
            specific_month=specific_month,
            force=force
        )
        
        # Phase 1: Process usage stage months
        console.print("🔄 Phase 1: Processing usage stage data...")
        usage_processed = process_usage_stage_months(
            config.dataset_name, 
            data_service, 
            specific_month
        )
        
        # Phase 2: Gap analysis and metadata processing  
        console.print("🔄 Phase 2: Analyzing gaps and processing metadata stage...")
        metadata_available = data_service.discover_available_months(config.dataset_name, "metadata")
        
        # Filter metadata stage to only providers with active processors
        active_providers = set(ProcessorFactory.ACTIVE_PROCESSORS.keys())
        skipped_providers_metadata = set(metadata_available.keys()) - active_providers
        metadata_available = {
            provider: months for provider, months in metadata_available.items() 
            if provider in active_providers
        }
        
        if skipped_providers_metadata:
            logger.info(
                "Skipping providers without active discount processors in metadata stage",
                dataset_name=config.dataset_name,
                skipped_providers=list(skipped_providers_metadata),
                active_providers=list(active_providers)
            )
        
        gaps = gap_analyzer.find_processing_gaps(
            config.dataset_name,
            usage_processed,
            metadata_available
        )
        
        metadata_processed = 0
        if gaps:
            console.print(f"📊 Found {len(gaps)} months to process from metadata stage")
            metadata_processed = process_metadata_gaps(config.dataset_name, gaps, data_service)
        else:
            console.print("✅ No gaps found - all providers processed from usage stage")
        
        # Write success marker
        total_processed = sum(len(months) for months in usage_processed.values()) + metadata_processed
        success_content = {
            "status": "complete",
            "dataset_name": config.dataset_name,
            "usage_months": sum(len(months) for months in usage_processed.values()),
            "metadata_gaps": metadata_processed,
            "total_months": total_processed,
            "force_used": force
        }
        
        data_service.write_success_marker(config.dataset_name, success_content)
        
        logger.info(
            "Discount allocation completed successfully",
            dataset_name=config.dataset_name,
            total_processed=total_processed,
            usage_months=success_content["usage_months"],
            metadata_gaps=metadata_processed
        )
        
        return ProcessingResult(
            success=True,
            processed_months=total_processed,
            usage_months=success_content["usage_months"],
            metadata_gaps=metadata_processed
        )
        
    except Exception as e:
        logger.error(
            "Discount allocation failed",
            dataset_name=config.dataset_name,
            error=str(e),
            error_type=type(e).__name__
        )
        return ProcessingResult(
            success=False,
            error=str(e)
        )


def process_usage_stage_months(
    dataset_name: str,
    data_service: DiscountDataService,
    specific_month: Optional[str] = None
) -> Dict[str, List[str]]:
    """
    Process all available months from usage stage for providers with active processors.
    
    Args:
        dataset_name: Dataset identifier
        data_service: S3 data service instance
        specific_month: Optional specific month filter
        
    Returns:
        Dict[provider, List[months]] that were successfully processed
    """
    logger = get_logger(__name__)
    
    # Discover available months from usage stage
    usage_available = data_service.discover_available_months(dataset_name, "usage")
    
    if not usage_available:
        logger.warning(
            "No months found in usage stage",
            dataset_name=dataset_name
        )
        return {}
    
    # Filter to only providers with active processors
    active_providers = set(ProcessorFactory.ACTIVE_PROCESSORS.keys())
    skipped_providers = set(usage_available.keys()) - active_providers
    usage_available = {
        provider: months for provider, months in usage_available.items() 
        if provider in active_providers
    }
    
    if skipped_providers:
        logger.info(
            "Skipping providers without active discount processors",
            dataset_name=dataset_name,
            skipped_providers=list(skipped_providers),
            active_providers=list(active_providers)
        )
    
    if not usage_available:
        logger.info(
            "No providers with active discount processors found in usage stage",
            dataset_name=dataset_name,
            active_providers=list(active_providers)
        )
        return {}
    
    # Filter by specific month if requested
    if specific_month:
        filtered_usage = {}
        for provider, months in usage_available.items():
            if specific_month in months:
                filtered_usage[provider] = [specific_month]
        usage_available = filtered_usage
    
    processed_successfully = {}
    
    # Process each provider/month combination
    for provider, months in usage_available.items():
        processed_successfully[provider] = []
        
        for month in months:
            console.print(f"🔄 Processing {provider} {month} (from usage stage)...")
            
            try:
                # TODO: Implement actual discount processing
                # For now, just simulate successful processing
                success = process_single_month(dataset_name, provider, month, "usage", data_service)
                
                if success:
                    processed_successfully[provider].append(month)
                    console.print(f"✅ {provider} {month}: Usage stage processing complete")
                else:
                    console.print(f"⚠️  {provider} {month}: Usage stage processing failed")
                    console.print(f"🛑 [red]PROCESSING STOPPED: Failure in {provider}/{month}[/red]")
                    console.print(f"💡 [yellow]Fix the error above before continuing with other months[/yellow]")
                    # Fail fast - stop processing immediately on first failure
                    return processed_successfully
                    
            except Exception as e:
                logger.error(
                    "Failed to process usage stage month",
                    dataset_name=dataset_name,
                    provider=provider,
                    month=month,
                    error=str(e)
                )
                console.print(f"❌ {provider} {month}: Error processing usage stage")
                console.print(f"🛑 [red]PROCESSING STOPPED: Exception in {provider}/{month}[/red]")
                console.print(f"💡 [yellow]Error: {str(e)}[/yellow]")
                # Fail fast - stop processing immediately on exception
                return processed_successfully
    
    # Remove providers with no successful months
    processed_successfully = {p: months for p, months in processed_successfully.items() if months}
    
    return processed_successfully


def process_metadata_gaps(
    dataset_name: str,
    gaps: List[ProcessingGap],
    data_service: DiscountDataService
) -> int:
    """
    Process missing months from metadata stage.
    
    Args:
        dataset_name: Dataset identifier
        gaps: List of processing gaps to fill
        data_service: S3 data service instance
        
    Returns:
        Number of months successfully processed
    """
    logger = get_logger(__name__)
    processed_count = 0
    
    for gap in gaps:
        console.print(f"🔄 Processing {gap.provider} {gap.month} (from metadata stage)...")
        
        try:
            # TODO: Implement actual discount processing with schema harmonization
            # For now, just simulate successful processing
            success = process_single_month(dataset_name, gap.provider, gap.month, "metadata", data_service)
            
            if success:
                processed_count += 1
                console.print(f"✅ {gap.provider} {gap.month}: Metadata gap processing complete")
            else:
                console.print(f"⚠️  {gap.provider} {gap.month}: Metadata gap processing failed")
                console.print(f"🛑 [red]PROCESSING STOPPED: Failure in {gap.provider}/{gap.month}[/red]")
                console.print(f"💡 [yellow]Fix the error above before continuing with other months[/yellow]")
                # Fail fast - stop processing immediately on first failure
                return processed_count
                
        except Exception as e:
            logger.error(
                "Failed to process metadata gap",
                dataset_name=dataset_name,
                provider=gap.provider,
                month=gap.month,
                error=str(e)
            )
            console.print(f"❌ {gap.provider} {gap.month}: Error processing metadata gap")
            console.print(f"🛑 [red]PROCESSING STOPPED: Exception in {gap.provider}/{gap.month}[/red]")
            console.print(f"💡 [yellow]Error: {str(e)}[/yellow]")
            # Fail fast - stop processing immediately on exception
            return processed_count
    
    return processed_count


def process_single_month(
    dataset_name: str,
    provider: str,
    month: str,
    source: str,
    data_service: DiscountDataService
) -> bool:
    """
    Process a single month from the specified source with full discount logic.
    
    Args:
        dataset_name: Dataset identifier
        provider: Cloud provider
        month: Month in YYYY-MM format
        source: Source stage ("usage" or "metadata")
        data_service: S3 data service instance
        
    Returns:
        True if processing successful
    """
    logger = get_logger(__name__)
    schema_harmonizer = SchemaHarmonizer()
    
    try:
        # Download data from appropriate source
        if source == "usage":
            data = data_service.download_usage_month_data(dataset_name, provider, month)
        else:
            data = data_service.download_metadata_month_data(dataset_name, provider, month)
        
        if data is None:
            logger.warning(
                f"No data found for {provider} {month} from {source} stage",
                dataset_name=dataset_name
            )
            return False
        
        logger.info(
            f"Successfully downloaded data for {provider} {month} from {source} stage",
            dataset_name=dataset_name,
            rows=len(data),
            columns=len(data.columns)
        )
        
        # Schema harmonization for metadata stage data
        if source == "metadata":
            data = schema_harmonizer.harmonize_metadata_schema(data)
            
            # Validate harmonization (optional, for debugging)
            if not schema_harmonizer.validate_harmonized_schema(data, data):
                logger.warning(f"Schema harmonization validation failed for {provider} {month}")
        
        # Ensure discount columns exist (in case they're missing)
        data = schema_harmonizer.ensure_discount_columns(data)
        
        # Apply discount processors for this provider
        enhanced_data, already_discounted = apply_discount_processors(
            data, provider, already_discounted=set()
        )
        
        # Upload to discount stage
        success = data_service.upload_discount_month_data(dataset_name, provider, month, enhanced_data)
        
        if success:
            logger.info(
                f"Successfully processed discount allocation for {provider} {month}",
                dataset_name=dataset_name,
                source=source,
                rows_processed=len(enhanced_data),
                resources_discounted=len(already_discounted)
            )
        
        return success
        
    except Exception as e:
        logger.error(
            f"Failed to process discount allocation for {provider} {month}",
            dataset_name=dataset_name,
            source=source,
            error=str(e),
            error_type=type(e).__name__,
            error_details=repr(e)
        )
        return False