"""Dual-source S3 data service for discount allocation."""

from typing import Dict, List, Optional
import polars as pl
from pathlib import Path

from src.services.s3_data_service import S3DataService
from src.services.environment_service import env_service
from src.log_context.core import get_logger


class DiscountDataService:
    """S3 service for discount stage data management with dual-source support.
    
    Follows the established S3 path structure with dual input sources:
    s3://bucket/czbillgen/datasets/{dataset_name}/
    ├── usage/{provider}/{month}.parquet      # INPUT (Phase 1)
    ├── metadata/{provider}/{month}.parquet   # INPUT (Phase 2 - gaps)
    └── discount/{provider}/{month}.parquet   # OUTPUT
    """
    
    def __init__(self):
        self.logger = get_logger(__name__)
        s3_config = env_service.get_s3_config()
        self.bucket = s3_config["bucket"]
        self.s3_service = S3DataService()
        
    def download_usage_month_data(self, dataset_name: str, provider: str, month: str) -> Optional[pl.DataFrame]:
        """
        Download month data from usage stage S3.
        
        Args:
            dataset_name: Dataset identifier
            provider: Cloud provider
            month: Month in YYYY-MM format
            
        Returns:
            DataFrame with billing data from usage stage, None if not found
        """
        try:
            # Use existing S3DataService pattern for consistency
            self.logger.debug(
                "Attempting to download usage stage data",
                dataset_name=dataset_name,
                provider=provider,
                month=month
            )
            
            data = self.s3_service._load_upload_month_data(dataset_name, "usage", provider, month)
            
            if data is None:
                self.logger.warning(
                    "Usage stage data returned None",
                    dataset_name=dataset_name,
                    provider=provider,
                    month=month
                )
                return None
            
            self.logger.info(
                "Downloaded month data from usage stage",
                dataset_name=dataset_name,
                provider=provider,
                month=month,
                rows=len(data)
            )
            
            return data
            
        except Exception as e:
            self.logger.error(
                "Failed to download usage stage data",
                dataset_name=dataset_name,
                provider=provider,
                month=month,
                error=str(e),
                error_type=type(e).__name__
            )
            return None
    
    def download_metadata_month_data(self, dataset_name: str, provider: str, month: str) -> Optional[pl.DataFrame]:
        """
        Download month data from metadata stage S3.
        
        Args:
            dataset_name: Dataset identifier
            provider: Cloud provider
            month: Month in YYYY-MM format
            
        Returns:
            DataFrame with billing data from metadata stage, None if not found
        """
        try:
            # Use existing S3DataService pattern for consistency
            data = self.s3_service._load_upload_month_data(dataset_name, "metadata", provider, month)
            
            self.logger.info(
                "Downloaded month data from metadata stage",
                dataset_name=dataset_name,
                provider=provider,
                month=month,
                rows=len(data)
            )
            
            return data
            
        except Exception as e:
            self.logger.error(
                "Failed to download month data from metadata stage",
                dataset_name=dataset_name,
                provider=provider,
                month=month,
                error=str(e)
            )
            return None
    
    def upload_discount_month_data(self, dataset_name: str, provider: str, month: str, data: pl.DataFrame) -> bool:
        """
        Upload enhanced month data to discount stage S3.
        
        Args:
            dataset_name: Dataset identifier
            provider: Cloud provider  
            month: Month in YYYY-MM format
            data: Enhanced DataFrame with discount columns
            
        Returns:
            True if upload successful
            
        Raises:
            Exception: If upload fails
        """
        try:
            # Use existing S3DataService upload method
            self.s3_service.upload_dataset_to_stage(dataset_name, "discount", provider, month, data)
            
            self.logger.info(
                "Uploaded month data to discount stage",
                dataset_name=dataset_name,
                provider=provider,
                month=month,
                rows=len(data)
            )
            
            return True
            
        except Exception as e:
            self.logger.error(
                "Failed to upload month data to discount stage",
                dataset_name=dataset_name,
                provider=provider,
                month=month,
                error=str(e)
            )
            raise
    
    def discover_available_months(self, dataset_name: str, stage: str) -> Dict[str, List[str]]:
        """
        Discover available months from usage or metadata stage by provider.
        
        Args:
            dataset_name: Dataset identifier
            stage: Stage name ("usage" or "metadata")
            
        Returns:
            Dict mapping provider to list of available months
        """
        try:
            # Get list of providers for this stage
            providers = self.s3_service.list_providers(dataset_name, stage)
            months_by_provider = {}
            
            for provider in providers:
                # Get months for each provider
                months = self.s3_service.list_months_fast(dataset_name, stage, provider)
                if months:
                    months_by_provider[provider] = months
            
            total_months = sum(len(months) for months in months_by_provider.values())
            self.logger.info(
                f"Discovered available months from {stage} stage",
                dataset_name=dataset_name,
                stage=stage,
                providers=list(months_by_provider.keys()),
                total_months=total_months
            )
            
            return months_by_provider
            
        except Exception as e:
            self.logger.error(
                f"Failed to discover months from {stage} stage",
                dataset_name=dataset_name,
                stage=stage,
                error=str(e)
            )
            return {}
    
    def write_success_marker(self, dataset_name: str, success_content: dict) -> None:
        """
        Write SUCCESS marker after all months processed.
        
        Args:
            dataset_name: Dataset identifier
            success_content: Content to write to SUCCESS marker
        """
        try:
            # Use existing success marker method from S3DataService
            self.s3_service.write_success_marker(dataset_name, success_content)
            
            self.logger.info(
                "Wrote SUCCESS marker for discount stage",
                dataset_name=dataset_name,
                content=success_content
            )
            
        except Exception as e:
            self.logger.error(
                "Failed to write SUCCESS marker for discount stage",
                dataset_name=dataset_name,
                error=str(e)
            )
            raise