"""
Metadata Data Service for S3-backed metadata allocation data management.

Provides S3 operations for metadata allocation stage:
- Reads from k8s stage S3 output (fallback to generate stage)
- Writes to metadata 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, Tuple

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
from ..provider_normalizer import ProviderNormalizer


class MetadataDataService:
    """
    Service for S3-backed metadata allocation data operations.

    Handles month-by-month processing with:
    - Reading from k8s stage S3 output (fallback to generate stage)
    - Writing to metadata stage S3 output
    - Local caching for memory efficiency
    - Success marker management
    """

    def __init__(self):
        """Initialize with S3 service and local cache."""
        self.logger = get_logger(__name__)
        self.s3_service = S3DataService()

        # Local cache configuration
        self.cache_dir = Path.home() / ".czbg" / "cache" / "metadata_data"
        self.cache_dir.mkdir(parents=True, exist_ok=True)

    def _normalize_provider_name(self, provider: str) -> str:
        """Normalize provider name to match config expectations."""
        return ProviderNormalizer.normalize(provider)

    def list_available_months(self, dataset_name: str, provider: str = None) -> List[Tuple[str, str]]:
        """
        List available provider-month combinations from S3.

        Checks k8s stage first, then falls back to generate stage.
        Returns original S3 provider names to maintain consistency with generate/k8s stages.

        Args:
            dataset_name: Dataset name
            provider: Optional provider filter (will be normalized for matching)

        Returns:
            List of (s3_provider, month) tuples available for processing
        """
        try:
            provider_months = set()

            # Get all providers if not specified
            if provider:
                # Find S3 provider that matches the normalized input
                k8s_providers = self.s3_service.list_providers(dataset_name, "k8s")
                generate_providers = self.s3_service.list_providers(dataset_name, "generate")
                all_providers = list(set(k8s_providers + generate_providers))

                providers_to_check = []
                for s3_provider in all_providers:
                    if self._normalize_provider_name(s3_provider) == self._normalize_provider_name(provider):
                        providers_to_check.append(s3_provider)
            else:
                # Try both k8s and generate stages to find all available providers
                k8s_providers = self.s3_service.list_providers(dataset_name, "k8s")
                generate_providers = self.s3_service.list_providers(dataset_name, "generate")
                providers_to_check = list(set(k8s_providers + generate_providers))

            # Check each provider for available months - preserve original S3 names
            for s3_provider in providers_to_check:
                # Try k8s stage first
                try:
                    k8s_months = self.s3_service.list_months_fast(dataset_name, "k8s", s3_provider)
                    for month in k8s_months:
                        provider_months.add((s3_provider, month))
                except ClientError as e:
                    if e.response['Error']['Code'] in ['NoSuchKey', 'NoSuchBucket']:
                        pass  # K8s stage might not exist for this provider
                    else:
                        raise  # Re-raise other S3 errors (including credential issues)

                # Try generate stage as fallback
                try:
                    generate_months = self.s3_service.list_months_fast(dataset_name, "generate", s3_provider)
                    for month in generate_months:
                        provider_months.add((s3_provider, month))
                except ClientError as e:
                    if e.response['Error']['Code'] in ['NoSuchKey', 'NoSuchBucket']:
                        pass  # Provider might not exist in generate either
                    else:
                        raise  # Re-raise other S3 errors (including credential issues)

            # Sort by provider, then month
            return sorted(list(provider_months))

        except (NoCredentialsError, ClientError) as e:
            # Let AWS credential and permission errors propagate with clear context
            self.logger.error(f"AWS access error: {e}")
            raise

    def _check_data_availability(self, dataset_name: str, provider: str, month: str, stage: str) -> bool:
        """
        Lightweight check for data availability in specific stage without downloading.

        Args:
            dataset_name: Dataset name
            provider: Provider name (will be normalized)
            month: Month string
            stage: Stage name ("k8s" or "generate")

        Returns:
            True if data exists in the specified stage
        """
        try:
            # Find S3 provider that matches normalized input
            stage_providers = self.s3_service.list_providers(dataset_name, stage)
            for s3_provider in stage_providers:
                if self._normalize_provider_name(s3_provider) == self._normalize_provider_name(provider):
                    # Check if month exists for this provider
                    months = self.s3_service.list_months_fast(dataset_name, stage, s3_provider)
                    return month in months

            return False

        except ClientError as e:
            if e.response['Error']['Code'] in ['NoSuchKey', 'NoSuchBucket']:
                return False  # Data doesn't exist
            else:
                # Log other S3 errors but don't fail initialization
                self.logger.debug(f"S3 error checking {stage} stage for {provider}/{month}: {e}")
                return False
        except Exception as e:
            self.logger.debug(f"Error checking {stage} stage availability for {provider}/{month}: {e}")
            return False

    def download_month_data(self, dataset_name: str, provider: str, month: str) -> Optional[pl.DataFrame]:
        """
        Download month data from S3 with k8s -> generate stage fallback.

        Tries k8s stage first, falls back to generate stage if not available.

        Args:
            dataset_name: Dataset name
            provider: Provider name (original S3 name, e.g. 'AWS')
            month: Month string (YYYY-MM)

        Returns:
            DataFrame with month data, or None if not found
        """
        # Use normalized provider name for cache file naming
        normalized_provider = self._normalize_provider_name(provider)
        cache_file = self.cache_dir / f"{dataset_name}_{normalized_provider}_{month}.parquet"

        try:
            # Try to load from local cache first
            if cache_file.exists():
                self.logger.debug(f"Loading {provider} {month} from cache")
                return pl.read_parquet(cache_file)

            # Use the provided provider name directly (it's already the S3 name)
            s3_provider = provider
            df = None
            data_source = None

            # Try k8s stage first
            k8s_key = f"czbillgen/datasets/{dataset_name}/k8s/{s3_provider}/{month}.parquet"
            if self.s3_service.file_exists(k8s_key):
                try:
                    df = self.s3_service.read_parquet_from_s3(k8s_key)
                    data_source = "k8s"
                    self.logger.debug(f"Downloaded {provider} {month} from k8s stage")
                except ClientError as e:
                    if e.response['Error']['Code'] in ['NoSuchKey', 'AccessDenied']:
                        self.logger.debug(f"K8s stage read failed for {provider} {month}: {e}")
                    else:
                        raise  # Re-raise credential and other serious S3 errors

            # Fallback to generate stage
            if df is None:
                generate_key = f"czbillgen/datasets/{dataset_name}/generate/{s3_provider}/{month}.parquet"
                if self.s3_service.file_exists(generate_key):
                    try:
                        df = self.s3_service.read_parquet_from_s3(generate_key)
                        data_source = "generate"
                        self.logger.debug(f"Downloaded {provider} {month} from generate stage")
                    except ClientError as e:
                        if e.response['Error']['Code'] in ['NoSuchKey', 'AccessDenied']:
                            self.logger.debug(f"Generate stage read failed for {provider} {month}: {e}")
                        else:
                            raise  # Re-raise credential and other serious S3 errors

            if df is None:
                self.logger.warning(f"No data found for {provider} {month} in k8s or generate stages")
                return None

            # Cache the downloaded data
            df.write_parquet(cache_file)

            self.logger.info(f"Downloaded {provider} {month} from {data_source} stage: {len(df):,} rows")
            return df

        except (NoCredentialsError, ClientError) as e:
            # Let AWS credential and permission errors propagate with clear context
            self.logger.error(f"AWS access error while downloading {provider} {month}: {e}")
            raise

    def upload_month_data(self, dataset_name: str, provider: str, month: str,
                         data: pl.DataFrame) -> bool:
        """
        Upload processed month data to metadata stage S3 output.

        Args:
            dataset_name: Dataset name
            provider: Provider name (original S3 name, e.g. 'AWS')
            month: Month string (YYYY-MM)
            data: Processed DataFrame

        Returns:
            True if upload successful, False otherwise
        """
        try:
            # Use the provided provider name directly (it's already the S3 name)
            s3_provider = provider

            # Upload the main data using S3DataService
            self.s3_service.upload_dataset_to_stage(
                dataset_name, "metadata", s3_provider, month, data
            )

            # Clean up local cache - use normalized name for cache file
            normalized_provider = self._normalize_provider_name(provider)
            cache_file = self.cache_dir / f"{dataset_name}_{normalized_provider}_{month}.parquet"
            if cache_file.exists():
                cache_file.unlink()

            self.logger.info(f"Successfully uploaded {s3_provider} {month}: {len(data):,} rows")
            return True

        except (NoCredentialsError, ClientError) as e:
            # Let AWS credential and permission errors propagate with clear context
            self.logger.error(f"AWS access error while uploading {provider} {month}: {e}")
            raise

    def check_completion_status(self, dataset_name: str, provider: str, month: str) -> bool:
        """
        Check if provider-month has been successfully processed by looking for the parquet file.

        Args:
            dataset_name: Dataset name
            provider: Provider name (original S3 name, e.g. 'AWS')
            month: Month string (YYYY-MM)

        Returns:
            True if processing completed successfully (parquet file exists)
        """
        try:
            # Check for the actual parquet file - same pattern as generate/k8s stages
            parquet_key = f"czbillgen/datasets/{dataset_name}/metadata/{provider}/{month}.parquet"
            return self.s3_service.file_exists(parquet_key)

        except (NoCredentialsError, ClientError) as e:
            # Let AWS credential and permission errors propagate with clear context
            self.logger.error(f"AWS access error while checking completion status for {provider} {month}: {e}")
            raise

    def write_stage_success_marker(self, dataset_name: str, processing_summary: Dict[str, Any]) -> bool:
        """
        Write stage-level success marker following generate/k8s pattern.

        Args:
            dataset_name: Dataset name
            processing_summary: Summary of all processing results

        Returns:
            True if success marker written successfully
        """
        try:
            # Create success marker content following standard pattern
            success_content = {
                "timestamp": datetime.utcnow().isoformat(),
                "dataset_name": dataset_name,
                "stage": "metadata",
                "status": "completed",
                "processing_summary": processing_summary
            }

            # Write to S3 using standard stage-level path
            success_key = f"czbillgen/datasets/{dataset_name}/metadata/_SUCCESS"
            self.s3_service.s3_client.put_object(
                Bucket=self.s3_service.bucket,
                Key=success_key,
                Body=json.dumps(success_content, indent=2),
                ContentType='application/json'
            )

            self.logger.info(f"Stage success marker written for dataset {dataset_name}")
            return True

        except (NoCredentialsError, ClientError) as e:
            # Let AWS credential and permission errors propagate with clear context
            self.logger.error(f"AWS access error while writing stage success marker: {e}")
            raise

    def cleanup_cache(self, dataset_name: str = None):
        """
        Clean up local cache files.

        Args:
            dataset_name: If provided, only clean cache for this dataset
        """
        try:
            if dataset_name:
                pattern = f"{dataset_name}_*.parquet"
                for cache_file in self.cache_dir.glob(pattern):
                    cache_file.unlink()
                    self.logger.debug(f"Cleaned cache file: {cache_file}")
            else:
                for cache_file in self.cache_dir.glob("*.parquet"):
                    cache_file.unlink()
                    self.logger.debug(f"Cleaned cache file: {cache_file}")

        except Exception as e:
            self.logger.warning(f"Error cleaning cache: {e}")