"""
Stage discovery service for automatic stage selection with lazy evaluation.

Implements the lazy evaluation approach where each provider is automatically
assigned the highest priority stage that contains complete data.

Priority order: discount > usage > metadata > k8s > generate
Data consistency: One stage per provider, fail fast if partial data found.
"""

from typing import Dict, List, Optional, Set
from dataclasses import dataclass
import logging

from ...services.s3_data_service import S3DataService


logger = logging.getLogger(__name__)


@dataclass
class StageDiscoveryResult:
    """Result of stage discovery for a single provider."""
    provider: str
    stage: str
    available_months: List[str]
    discovery_method: str  # "discovered" or "override"
    

class StageDiscoveryError(Exception):
    """Raised when stage discovery fails."""
    
    def __init__(self, provider: str, dataset: str, message: str):
        self.provider = provider
        self.dataset = dataset
        super().__init__(f"Stage discovery failed for {provider} in dataset {dataset}: {message}")


class PartialDataError(StageDiscoveryError):
    """Raised when a stage has partial data (some but not all required months)."""
    
    def __init__(self, provider: str, dataset: str, stage: str, available_months: List[str], missing_months: List[str]):
        self.stage = stage
        self.available_months = available_months
        self.missing_months = missing_months
        
        message = (
            f"Provider {provider} in {stage} stage has partial data. "
            f"Available: {sorted(available_months)}, Missing: {sorted(missing_months)}. "
            f"Cannot mix stages - fix {stage} stage or exclude missing months from config."
        )
        super().__init__(provider, dataset, message)


class NoDataError(StageDiscoveryError):
    """Raised when no stage contains any data for a provider."""
    
    def __init__(self, provider: str, dataset: str, checked_stages: List[str]):
        self.checked_stages = checked_stages
        
        stage_status = ", ".join([f"{stage} (0 files)" for stage in checked_stages])
        message = f"No data found in any stage. Checked stages: {stage_status}"
        super().__init__(provider, dataset, message)


class StageDiscoveryService:
    """
    Service for automatic stage discovery with lazy evaluation.
    
    Implements the complete stage discovery logic:
    - Priority-based stage selection (discount > usage > metadata > k8s > generate)
    - Complete data validation (all required months must be in chosen stage)
    - Fail-fast on partial data (cannot mix stages for a provider)
    - Override support for manual stage specification
    """
    
    # Stage priority order (highest to lowest priority)
    STAGE_PRIORITY = ["discount", "usage", "metadata", "k8s", "generate"]
    
    def __init__(self):
        """Initialize stage discovery service."""
        self.s3_data_service = S3DataService()
        self._cache = {}  # Cache discovery results to avoid repeated S3 calls
    
    def discover_provider_stages(
        self, 
        dataset_name: str, 
        providers: List[str], 
        required_months: List[str],
        stage_overrides: Optional[Dict[str, str]] = None
    ) -> List[StageDiscoveryResult]:
        """
        Discover stages for multiple providers with complete data validation.
        
        Args:
            dataset_name: Dataset name to search in
            providers: List of provider names to discover stages for
            required_months: List of months that must be present in chosen stage
            stage_overrides: Optional dict mapping provider to forced stage
            
        Returns:
            List of StageDiscoveryResult objects with discovery results
            
        Raises:
            StageDiscoveryError: If discovery fails for any provider
            PartialDataError: If any stage has partial data for a provider
            NoDataError: If no stage has data for a provider
        """
        if not providers:
            return []
            
        if not required_months:
            raise ValueError("Required months list cannot be empty")
        
        stage_overrides = stage_overrides or {}
        results = []
        
        logger.info(f"Discovering stages for {len(providers)} providers in dataset {dataset_name}")
        logger.info(f"Required months: {required_months}")
        if stage_overrides:
            logger.info(f"Stage overrides: {stage_overrides}")
        
        for provider in providers:
            try:
                # Check for override first
                if provider in stage_overrides:
                    result = self._validate_override_stage(
                        dataset_name, provider, required_months, stage_overrides[provider]
                    )
                else:
                    # Use discovery
                    result = self._discover_provider_stage(
                        dataset_name, provider, required_months
                    )
                
                results.append(result)
                logger.info(f"✅ {provider}: {result.stage} stage ({result.discovery_method} - {len(result.available_months)} months)")
                
            except StageDiscoveryError as e:
                logger.error(f"❌ {provider}: {e}")
                raise  # Re-raise to stop processing
        
        return results
    
    def _discover_provider_stage(
        self, 
        dataset_name: str, 
        provider: str, 
        required_months: List[str]
    ) -> StageDiscoveryResult:
        """
        Discover highest priority stage with complete data for provider.
        
        Args:
            dataset_name: Dataset name to search in
            provider: Provider name
            required_months: Months that must be present
            
        Returns:
            StageDiscoveryResult with discovered stage
            
        Raises:
            PartialDataError: If any stage has partial data
            NoDataError: If no stage has data
        """
        cache_key = f"{dataset_name}:{provider}:{':'.join(sorted(required_months))}"
        if cache_key in self._cache:
            return self._cache[cache_key]
        
        logger.debug(f"Discovering stage for {provider} with months {required_months}")
        
        required_months_set = set(required_months)
        
        for stage in self.STAGE_PRIORITY:
            available_months = self._get_provider_months(dataset_name, stage, provider)
            
            if len(available_months) == 0:
                # Stage has no data for provider - continue to next stage
                logger.debug(f"  {stage}: 0 months available - continuing")
                continue
                
            available_months_set = set(available_months)
            
            if required_months_set.issubset(available_months_set):
                # Stage has ALL required months - use this stage
                result = StageDiscoveryResult(
                    provider=provider,
                    stage=stage,
                    available_months=available_months,
                    discovery_method="discovered"
                )
                self._cache[cache_key] = result
                logger.debug(f"  {stage}: {len(available_months)} months available - SELECTED")
                return result
                
            else:
                # Stage has SOME but not ALL months - FAIL FAST
                missing_months = list(required_months_set - available_months_set)
                logger.debug(f"  {stage}: Partial data - available: {available_months}, missing: {missing_months}")
                raise PartialDataError(
                    provider, dataset_name, stage, available_months, missing_months
                )
        
        # No stage has any data
        raise NoDataError(provider, dataset_name, self.STAGE_PRIORITY)
    
    def _validate_override_stage(
        self, 
        dataset_name: str, 
        provider: str, 
        required_months: List[str], 
        override_stage: str
    ) -> StageDiscoveryResult:
        """
        Validate that override stage has complete data for provider.
        
        Args:
            dataset_name: Dataset name
            provider: Provider name
            required_months: Months that must be present
            override_stage: Stage specified in override
            
        Returns:
            StageDiscoveryResult with validated override
            
        Raises:
            ValueError: If override stage is invalid
            PartialDataError: If override stage has partial data
            NoDataError: If override stage has no data
        """
        if override_stage not in self.STAGE_PRIORITY:
            raise ValueError(f"Invalid override stage '{override_stage}' for provider {provider}. Valid stages: {self.STAGE_PRIORITY}")
        
        available_months = self._get_provider_months(dataset_name, override_stage, provider)
        required_months_set = set(required_months)
        available_months_set = set(available_months)
        
        if len(available_months) == 0:
            # Override stage has no data
            raise NoDataError(provider, dataset_name, [override_stage])
            
        elif required_months_set.issubset(available_months_set):
            # Override stage has all required months - valid
            return StageDiscoveryResult(
                provider=provider,
                stage=override_stage,
                available_months=available_months,
                discovery_method="override"
            )
            
        else:
            # Override stage has partial data - fail fast
            missing_months = list(required_months_set - available_months_set)
            raise PartialDataError(
                provider, dataset_name, override_stage, available_months, missing_months
            )
    
    def _get_provider_months(self, dataset_name: str, stage: str, provider: str) -> List[str]:
        """
        Get list of available months for provider in specific stage.
        
        Args:
            dataset_name: Dataset name
            stage: Stage name
            provider: Provider name
            
        Returns:
            List of available month strings (e.g., ["2025-01", "2025-02"])
        """
        try:
            # Use existing S3DataService method to get fast summary
            summary = self.s3_data_service.get_upload_dataset_summary_fast(dataset_name, stage)
            
            if provider in summary:
                return sorted(summary[provider])  # Return sorted list of months
            else:
                return []  # Provider not in this stage
                
        except Exception as e:
            # Stage doesn't exist or other error
            logger.debug(f"Error getting months for {provider} in {stage} stage: {e}")
            return []
    
    def format_discovery_results(self, results: List[StageDiscoveryResult]) -> str:
        """
        Format discovery results for user-friendly display.
        
        Args:
            results: List of discovery results
            
        Returns:
            Formatted string for console output
        """
        if not results:
            return "No discovery results to display"
        
        lines = ["🔍 Stage Discovery Results:"]
        
        discovery_count = 0
        override_count = 0
        
        for result in results:
            method_desc = f"({result.discovery_method} - {len(result.available_months)} months available)"
            lines.append(f"   {result.provider}: {result.stage} {method_desc}")
            
            if result.discovery_method == "discovered":
                discovery_count += 1
            else:
                override_count += 1
        
        # Add summary
        lines.append("")
        lines.append("ℹ️  Discovery Summary:")
        if override_count > 0:
            lines.append(f"   • {override_count} providers using overrides")
        if discovery_count > 0:
            lines.append(f"   • {discovery_count} providers auto-discovered")
        
        # Check for mixed stages
        stages_used = set(result.stage for result in results)
        if len(stages_used) > 1:
            lines.append("   • Mixed stages detected - upload will proceed with best available data per provider")
        
        lines.append("")
        lines.append("✅ All providers have complete data - ready to upload")
        
        return "\n".join(lines)