"""
K8s Data Service for S3-backed K8s allocation data management.

Provides S3 operations for K8s allocation stage:
- Reads from generate stage S3 output
- Writes to k8s stage S3 output
- Month-by-month processing with local caching
- Memory-efficient operations
"""

import json
import os
from datetime import datetime
from pathlib import Path
from typing import Optional, List, Dict, Any

import polars as pl
from botocore.exceptions import ClientError, NoCredentialsError

from src.log_context import get_logger
from src.services.s3_data_service import S3DataService


class K8sDataService:
    """
    Service for S3-backed K8s allocation data operations.
    
    Handles month-by-month processing with:
    - Reading from generate stage S3 output
    - Writing to k8s stage S3 output
    - Local caching for performance
    - Memory-efficient operations
    """
    
    def __init__(self):
        """Initialize K8s data service with S3 delegation."""
        self.logger = get_logger(__name__)
        
        # Delegate to S3DataService for all S3 operations
        self.s3_service = S3DataService()
        
        # Local cache configuration
        self.cache_dir = Path.home() / ".czbg" / "cache" / "k8s_data"
        self.cache_dir.mkdir(parents=True, exist_ok=True)
    
    def list_available_months(self, dataset_name: str, provider: str = None) -> List[str]:
        """
        List available months from generate stage S3 output.
        
        Args:
            dataset_name: Dataset name
            provider: Optional provider filter (AWS/GCP/Azure)
            
        Returns:
            List of month strings (YYYY-MM) available for processing
        """
        try:
            months = set()
            
            # Get all providers if not specified
            if provider:
                providers_to_check = [provider]
            else:
                providers_to_check = self.s3_service.list_providers(dataset_name, "generate")
            
            # Get months for each provider
            for provider_name in providers_to_check:
                provider_months = self.s3_service.list_months_fast(dataset_name, "generate", provider_name)
                months.update(provider_months)
            
            self.logger.info(
                "Dataset discovery completed",
                dataset_name=dataset_name,
                provider=provider,
                months_found=sorted(list(months))
            )
            
            return sorted(list(months))
            
        except Exception as e:
            self.logger.error(
                "Failed to list available months",
                dataset_name=dataset_name,
                provider=provider,
                error=str(e)
            )
            return []
    
    def download_month_data(self, dataset_name: str, provider: str, month: str) -> Optional[pl.DataFrame]:
        """
        Download month data from generate stage S3 output.
        
        Args:
            dataset_name: Dataset name
            provider: Provider name (AWS/GCP/Azure)
            month: Month string (YYYY-MM)
            
        Returns:
            DataFrame with month data or None if not found
        """
        try:
            # Construct S3 path following generate stage pattern
            s3_key = f"czbillgen/datasets/{dataset_name}/generate/{provider}/{month}.parquet"
            
            # Check if file exists using S3DataService
            if not self.s3_service.file_exists(s3_key):
                self.logger.warning(
                    "Month data not found",
                    dataset_name=dataset_name,
                    provider=provider,
                    month=month,
                    s3_key=s3_key
                )
                return None
            
            # Download to local cache
            local_path = self.cache_dir / f"{dataset_name}_{provider}_{month}.parquet"
            
            self.s3_service.s3_client.download_file(
                self.s3_service.bucket,
                s3_key,
                str(local_path)
            )
            
            # Read and return DataFrame
            df = pl.read_parquet(local_path)
            
            self.logger.info(
                "Month data downloaded",
                dataset_name=dataset_name,
                provider=provider,
                month=month,
                rows=df.height,
                local_path=str(local_path)
            )
            
            return df
            
        except Exception as e:
            self.logger.error(
                "Failed to download month data",
                dataset_name=dataset_name,
                provider=provider,
                month=month,
                error=str(e)
            )
            return None
    
    def upload_month_data(self, dataset_name: str, provider: str, month: str, 
                         data: pl.DataFrame, validation_report: Dict[str, Any] = None) -> bool:
        """
        Upload enhanced month data to k8s stage S3 output.
        
        Args:
            dataset_name: Dataset name
            provider: Provider name (AWS/GCP/Azure)
            month: Month string (YYYY-MM)
            data: Enhanced DataFrame with K8s metadata
            validation_report: Optional validation report (written to local cache only)
            
        Returns:
            True if successful, False otherwise
        """
        try:
            # Write to local cache first
            local_data_path = self.cache_dir / f"{dataset_name}_{provider}_{month}_k8s.parquet"
            data.write_parquet(local_data_path)
            
            # Upload data file using S3DataService
            self.s3_service.upload_dataset_to_stage(dataset_name, "k8s", provider, month, data)
            
            # Write validation report to local cache only (not uploaded to S3)
            if validation_report:
                local_report_path = self.cache_dir / f"{dataset_name}_{provider}_{month}_validation.json"
                with open(local_report_path, 'w') as f:
                    json.dump(validation_report, f, indent=2)
            
            self.logger.info(
                "Month data uploaded",
                dataset_name=dataset_name,
                provider=provider,
                month=month,
                rows=data.height,
                validation_report_cached=validation_report is not None
            )
            
            return True
            
        except Exception as e:
            self.logger.error(
                "Failed to upload month data",
                dataset_name=dataset_name,
                provider=provider,
                month=month,
                error=str(e)
            )
            return False
    
    def check_month_exists(self, dataset_name: str, provider: str, month: str) -> bool:
        """
        Check if month already processed in k8s stage.
        
        Args:
            dataset_name: Dataset name
            provider: Provider name (AWS/GCP/Azure)
            month: Month string (YYYY-MM)
            
        Returns:
            True if month already processed, False otherwise
        """
        try:
            s3_key = f"czbillgen/datasets/{dataset_name}/k8s/{provider}/{month}.parquet"
            
            return self.s3_service.file_exists(s3_key)
            
        except Exception as e:
            self.logger.error(
                "Failed to check month existence",
                dataset_name=dataset_name,
                provider=provider,
                month=month,
                error=str(e)
            )
            return False
    
    def write_success_marker(self, dataset_name: str, success_content: Dict[str, Any]) -> bool:
        """
        Write SUCCESS marker after all months processed.
        
        Args:
            dataset_name: Dataset name
            success_content: Success metadata
            
        Returns:
            True if successful, False otherwise
        """
        try:
            # Construct SUCCESS marker path
            success_key = f"czbillgen/datasets/{dataset_name}/k8s/_SUCCESS"
            
            # Write to local cache first
            local_success_path = self.cache_dir / f"{dataset_name}_SUCCESS.json"
            with open(local_success_path, 'w') as f:
                json.dump(success_content, f, indent=2)
            
            # Upload SUCCESS marker using S3DataService
            self.s3_service.upload_file(str(local_success_path), success_key)
            
            self.logger.info(
                "SUCCESS marker written",
                dataset_name=dataset_name,
                s3_key=success_key
            )
            
            return True
            
        except Exception as e:
            self.logger.error(
                "Failed to write SUCCESS marker",
                dataset_name=dataset_name,
                error=str(e)
            )
            return False
    
    def cleanup_local_cache(self, dataset_name: str = None, provider: str = None, month: str = None, validation_passed: bool = True):
        """
        Clean up local cache files.
        
        Args:
            dataset_name: Optional dataset filter
            provider: Optional provider filter
            month: Optional month filter
            validation_passed: If True, clean validation files; if False, preserve them for debugging
        """
        try:
            if dataset_name and provider and month:
                # Clean specific month
                patterns = [
                    f"{dataset_name}_{provider}_{month}.parquet",
                    f"{dataset_name}_{provider}_{month}_k8s.parquet"
                ]
                
                # Only clean validation files if validation passed
                if validation_passed:
                    patterns.append(f"{dataset_name}_{provider}_{month}_validation.json")
                
                for pattern in patterns:
                    file_path = self.cache_dir / pattern
                    if file_path.exists():
                        file_path.unlink()
                        
                self.logger.debug(
                    "Cleaned local cache for month",
                    dataset_name=dataset_name,
                    provider=provider,
                    month=month,
                    validation_files_cleaned=validation_passed
                )
            else:
                # Clean all cache files
                for file_path in self.cache_dir.glob("*"):
                    if file_path.is_file():
                        file_path.unlink()
                        
                self.logger.debug("Cleaned all local cache files")
                
        except Exception as e:
            self.logger.error(
                "Failed to cleanup local cache",
                error=str(e)
            )