"""S3 data service for usage allocation following established patterns."""

from typing import Dict, List, Optional, Tuple
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
from ..core.stage_discovery import StageDiscovery


class UsageDataService:
    """S3 service for usage stage data management.
    
    Follows the established S3 path structure:
    s3://bucket/czbillgen/datasets/{dataset_name}/
    ├── metadata/{provider}/{month}.parquet   # INPUT
    └── usage/{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()
        self.stage_discovery = StageDiscovery(self.s3_service)
        
    def download_month_data(
        self,
        dataset_name: str,
        provider: str,
        month: str,
        stage: str = "metadata"
    ) -> pl.DataFrame:
        """
        Download month data from specified pipeline stage.

        Args:
            dataset_name: Dataset identifier
            provider: Cloud provider
            month: Month in YYYY-MM format
            stage: Pipeline stage to load from (default: "metadata")

        Returns:
            DataFrame with billing data from specified stage

        Raises:
            FileNotFoundError: If stage data doesn't exist
        """
        try:
            # Use existing S3DataService pattern for consistency
            data = self.s3_service._load_upload_month_data(dataset_name, stage, provider, month)

            self.logger.info(
                f"Downloaded month data from {stage} stage",
                dataset_name=dataset_name,
                provider=provider,
                month=month,
                stage=stage,
                rows=len(data)
            )

            return data

        except Exception as e:
            self.logger.error(
                f"Failed to download month data from {stage} stage",
                dataset_name=dataset_name,
                provider=provider,
                month=month,
                stage=stage,
                error=str(e)
            )
            raise FileNotFoundError(f"{stage.title()} stage data not found for {provider} {month}") from e
    
    def upload_month_data(self, dataset_name: str, provider: str, month: str, data: pl.DataFrame) -> bool:
        """
        Upload enhanced month data to usage stage S3.
        
        Args:
            dataset_name: Dataset identifier
            provider: Cloud provider  
            month: Month in YYYY-MM format
            data: Enhanced DataFrame with usage columns
            
        Returns:
            True if upload successful
            
        Raises:
            Exception: If upload fails
        """
        try:
            # Use existing S3DataService pattern for consistency
            self.s3_service.upload_dataset_to_stage(dataset_name, "usage", provider, month, data)
            success = True
            
            if success:
                self.logger.info(
                    "Uploaded month data to usage stage",
                    dataset_name=dataset_name,
                    provider=provider,
                    month=month,
                    rows=len(data)
                )
            else:
                self.logger.error(
                    "Failed to upload month data to usage stage",
                    dataset_name=dataset_name,
                    provider=provider,
                    month=month
                )
            
            return success
            
        except Exception as e:
            self.logger.error(
                "Error uploading month data to usage stage",
                dataset_name=dataset_name,
                provider=provider,
                month=month,
                error=str(e)
            )
            raise
    
    def list_available_months(self, dataset_name: str, provider: Optional[str] = None) -> Dict[str, List[str]]:
        """
        List available months from metadata stage by provider.
        
        Args:
            dataset_name: Dataset identifier
            provider: Optional specific provider filter
            
        Returns:
            Dict[provider, List[months]] available for processing
        """
        try:
            # Use existing S3DataService pattern for consistency
            provider_months = self.s3_service.get_upload_dataset_summary_fast(dataset_name, "metadata")
            
            # Filter by provider if specified
            if provider:
                provider_months = {p: months for p, months in provider_months.items() if p == provider}
            
            self.logger.info(
                "Listed available months from metadata stage",
                dataset_name=dataset_name,
                providers=list(provider_months.keys()),
                total_months=sum(len(months) for months in provider_months.values())
            )
            
            return provider_months
            
        except Exception as e:
            self.logger.error(
                "Failed to list available months",
                dataset_name=dataset_name,
                provider=provider,
                error=str(e)
            )
            return {}

    def list_available_months_with_fallback(
        self,
        dataset_name: str,
        provider: Optional[str] = None
    ) -> Tuple[Dict[str, List[str]], Optional[str]]:
        """
        List available months with stage fallback.

        Tries stages in priority order: metadata → k8s → generate
        Returns first stage with data.

        Args:
            dataset_name: Dataset identifier
            provider: Optional specific provider filter

        Returns:
            Tuple of (months_dict, source_stage)
            - months_dict: Dict[provider, List[months]]
            - source_stage: Stage name where data was found, or None if no data
        """
        for stage in StageDiscovery.STAGE_PRIORITY:
            try:
                # Try to get months from this stage
                provider_months = self.s3_service.get_upload_dataset_summary_fast(
                    dataset_name, stage
                )

                # Found data in this stage
                if provider_months:
                    # Filter by provider if specified
                    if provider:
                        provider_months = {
                            p: months
                            for p, months in provider_months.items()
                            if p == provider
                        }

                    self.logger.info(
                        f"Found input data in {stage} stage",
                        dataset_name=dataset_name,
                        stage=stage,
                        providers=list(provider_months.keys()),
                        total_months=sum(len(months) for months in provider_months.values())
                    )

                    return provider_months, stage

            except Exception as e:
                # Stage not available - try next stage
                self.logger.debug(
                    f"Stage {stage} not available, trying next",
                    dataset_name=dataset_name,
                    stage=stage,
                    error=str(e)
                )
                continue

        # No data found in any stage
        self.logger.warning(
            "No input data found in any pipeline stage",
            dataset_name=dataset_name,
            checked_stages=StageDiscovery.STAGE_PRIORITY
        )
        return {}, None

    def file_exists(self, s3_path: str) -> bool:
        """
        Check if file exists in S3.
        
        Args:
            s3_path: S3 path to check
            
        Returns:
            True if file exists
        """
        try:
            # Use existing S3DataService pattern for consistency
            return self.s3_service.file_exists(s3_path)
                    
        except Exception as e:
            self.logger.error(f"Error checking file existence: {s3_path}", error=str(e))
            return False
    
    def delete_file(self, s3_path: str) -> bool:
        """
        Delete file from S3.
        
        Args:
            s3_path: S3 path to delete
            
        Returns:
            True if deletion successful
            
        Raises:
            Exception: If deletion fails
        """
        try:
            # Use existing S3DataService client for consistency
            self.s3_service.s3_client.delete_object(Bucket=self.bucket, Key=s3_path)
            self.logger.info(f"Deleted S3 file: {s3_path}")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to delete S3 file: {s3_path}", error=str(e))
            raise
    
    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: Success marker content
        """
        success_path = f"czbillgen/datasets/{dataset_name}/usage/_SUCCESS"
        
        try:
            import json
            from datetime import datetime
            
            # Add timestamp to success content
            success_content_with_timestamp = {
                **success_content,
                "completed_at": datetime.utcnow().isoformat() + "Z"
            }
            
            # Write JSON content to S3 using properly configured client
            json_content = json.dumps(success_content_with_timestamp, indent=2)
            
            self.s3_service.s3_client.put_object(
                Bucket=self.bucket,
                Key=success_path,
                Body=json_content.encode('utf-8'),
                ContentType='application/json'
            )
            
            self.logger.info(
                "Wrote usage stage SUCCESS marker",
                dataset_name=dataset_name,
                s3_path=success_path,
                content=success_content_with_timestamp
            )
            
        except Exception as e:
            self.logger.error(
                "Failed to write SUCCESS marker",
                dataset_name=dataset_name,
                error=str(e)
            )
            raise