from typing import List, Optional, Tuple
"""
ParquetDataService - Handle parquet operations for pipeline stages.

Primary Responsibility: Parquet file operations for generate, usage, discount, tax, k8s, metadata stages

Capabilities:
- Reference data loading with format detection
- Generated data loading from provider/month structure
- Append-only operations for immediate writes
- Schema validation before append operations
- Provider/month file organization
"""

from dataclasses import dataclass
from enum import Enum
from pathlib import Path

import polars as pl


class CompressionType(Enum):
    """Supported compression types for parquet files."""

    SNAPPY = "snappy"
    LZ4 = "lz4"
    NONE = None


@dataclass
class ParquetSaveOptions:
    """
    Options for saving parquet files.

    Configures compression and other save parameters for parquet operations.
    """

    compression: CompressionType = CompressionType.SNAPPY
    create_directories: bool = True

    @property
    def compression_value(self) -> Optional[str]:
        """Get compression value for polars."""
        return (
            self.compression.value if self.compression != CompressionType.NONE else None
        )


class ParquetDataService:
    """
    Service for parquet data operations in pipeline stages.

    Provides centralized parquet I/O operations for:
    - Reference data loading
    - Generated billing data operations
    - Provider/month file structure management
    - Append-only immediate writes

    Features:
    - Format detection and fallback mechanisms
    - Schema validation before append operations
    - Provider/month directory structure support
    - Memory-efficient append operations
    - Clear error reporting
    """

    def __init__(self):
        """Initialize ParquetDataService."""
        pass

    def load_reference_data(self, path: Path) -> pl.DataFrame:
        """
        Load reference data with format detection.

        Supports both parquet and CSV reference data files with
        automatic format detection and fallback mechanisms.

        Args:
            path: Path to reference data file

        Returns:
            Loaded reference DataFrame

        Raises:
            FileNotFoundError: If reference data file doesn't exist
            ValueError: If file cannot be parsed
        """
        if not path.exists():
            raise FileNotFoundError(
                f"Reference data file not found: {path}\n"
                f"Current working directory: {Path.cwd()}\n"
                f"Absolute path attempted: {path.resolve()}"
            )

        # Load based on file extension, fail fast on read errors
        if path.suffix.lower() == ".parquet":
            return pl.read_parquet(path)
        elif path.suffix.lower() == ".csv":
            return pl.read_csv(path)
        else:
            # Try parquet first for unknown extensions, then CSV
            try:
                return pl.read_parquet(path)
            except Exception:
                return pl.read_csv(path)

    def load_generated_data(
        self, base_dir: Path, provider: str, month: str
    ) -> pl.DataFrame:
        """
        Load generated data from provider/month structure.

        Loads data from the standard provider/month organization:
        {base_dir}/{provider}/{month}/data.parquet

        Args:
            base_dir: Base directory for provider/month structure
            provider: Provider name (e.g., "AWS", "GCP")
            month: Month in YYYY-MM format

        Returns:
            Loaded DataFrame

        Raises:
            FileNotFoundError: If data file doesn't exist
            ValueError: If file cannot be parsed
        """
        data_file = base_dir / provider / month / "data.parquet"

        if not data_file.exists():
            raise FileNotFoundError(
                f"Generated data file not found: {data_file}\n"
                f"Expected provider/month structure: {base_dir}/{provider}/{month}/data.parquet"
            )

        return pl.read_parquet(data_file)

    def append_to_provider_month_file(
        self,
        data: pl.DataFrame,
        base_dir: Path,
        provider: str,
        month: str,
        options: Optional[ParquetSaveOptions] = None,
    ) -> Path:
        """
        Append data to provider/month file with schema validation.

        This is the core function for immediate write pattern. Appends data
        to existing provider/month files or creates new ones if they don't exist.

        Performs schema validation to ensure consistency across appends.

        Args:
            data: DataFrame to append
            base_dir: Base directory for provider/month structure
            provider: Provider name (e.g., "AWS", "GCP")
            month: Month in YYYY-MM format
            options: Save options (compression, etc.)

        Returns:
            Path to the data file that was written

        Raises:
            ValueError: If data is empty or schema mismatch
        """
        if data is None or data.is_empty():
            raise ValueError("Cannot append empty DataFrame")

        if options is None:
            options = ParquetSaveOptions()

        # Validate month format
        if not self._is_valid_month_format(month):
            raise ValueError(f"Invalid month format '{month}'. Expected YYYY-MM.")

        # Create provider/month directory structure
        provider_month_dir = base_dir / provider / month
        if options.create_directories:
            provider_month_dir.mkdir(parents=True, exist_ok=True)

        # Determine data file path
        data_file = provider_month_dir / "data.parquet"

        # Handle append vs new file
        if data_file.exists():
            # Validate schema before append
            self._validate_schema_for_append(data, data_file)

            # Append to existing file
            # Read existing data, concatenate, and write back
            # This is the safest approach for schema consistency
            existing_data = pl.read_parquet(data_file)
            combined_data = pl.concat([existing_data, data], how="vertical")
            combined_data.write_parquet(
                data_file, compression=options.compression_value
            )
        else:
            # Create new file
            data.write_parquet(data_file, compression=options.compression_value)

        return data_file

    def save_data(
        self,
        data: pl.DataFrame,
        destination: Path,
        options: Optional[ParquetSaveOptions] = None,
    ) -> None:
        """
        Save DataFrame to parquet file.

        General-purpose parquet saving with compression and directory creation.

        Args:
            data: DataFrame to save
            destination: Target file path
            options: Save options (compression, etc.)

        Raises:
            ValueError: If data is empty or invalid
            PermissionError: If cannot write to destination
        """
        if data is None:
            raise ValueError("Cannot save None DataFrame")

        if not isinstance(data, pl.DataFrame):
            raise TypeError(f"Expected pl.DataFrame, got {type(data)}")

        if data.is_empty():
            raise ValueError("Cannot save empty DataFrame")

        if options is None:
            options = ParquetSaveOptions()

        # Create parent directories if needed
        if options.create_directories:
            destination.parent.mkdir(parents=True, exist_ok=True)

        data.write_parquet(destination, compression=options.compression_value)

    def get_provider_month_path(
        self, base_dir: Path, provider: str, month: str
    ) -> Path:
        """
        Get path to provider/month data file.

        Args:
            base_dir: Base directory for provider/month structure
            provider: Provider name
            month: Month in YYYY-MM format

        Returns:
            Path to data file (may not exist yet)
        """
        return base_dir / provider / month / "data.parquet"

    def list_provider_months(self, base_dir: Path) -> List[Tuple[str, str]]:
        """
        List all provider/month combinations in directory.

        Args:
            base_dir: Base directory to scan

        Returns:
            List of (provider, month) tuples
        """
        provider_months = []

        if not base_dir.exists():
            return provider_months

        for provider_dir in base_dir.iterdir():
            if provider_dir.is_dir() and provider_dir.name in ["AWS", "Azure", "GCP"]:
                for month_dir in provider_dir.iterdir():
                    if month_dir.is_dir() and self._is_valid_month_format(
                        month_dir.name
                    ):
                        data_file = month_dir / "data.parquet"
                        if data_file.exists():
                            provider_months.append((provider_dir.name, month_dir.name))

        return sorted(provider_months)

    def load_all_provider_month_data(self, base_dir: Path) -> pl.DataFrame:
        """
        Load and combine all provider/month data files.

        Args:
            base_dir: Base directory containing provider/month structure

        Returns:
            Combined DataFrame from all provider/month files

        Raises:
            FileNotFoundError: If no data files found
        """
        provider_months = self.list_provider_months(base_dir)

        if not provider_months:
            raise FileNotFoundError(f"No provider/month data files found in {base_dir}")

        dataframes = []
        for provider, month in provider_months:
            df = self.load_generated_data(base_dir, provider, month)
            dataframes.append(df)

        if not dataframes:
            raise FileNotFoundError(f"No valid data files found in {base_dir}")

        # Combine all dataframes
        return pl.concat(dataframes, how="vertical")

    def _validate_schema_for_append(
        self, new_data: pl.DataFrame, existing_file: Path
    ) -> None:
        """
        Validate that new data schema matches existing file schema.

        Args:
            new_data: DataFrame to append
            existing_file: Path to existing parquet file

        Raises:
            ValueError: If schema mismatch detected
        """
        # Use lazy frame to get schema without loading entire file
        existing_schema = pl.scan_parquet(existing_file).schema
        new_schema = new_data.schema

        # Check column names match
        existing_columns = set(existing_schema.keys())
        new_columns = set(new_schema.keys())

        if existing_columns != new_columns:
            missing_in_new = existing_columns - new_columns
            extra_in_new = new_columns - existing_columns

            error_parts = []
            if missing_in_new:
                error_parts.append(
                    f"Missing columns in new data: {sorted(missing_in_new)}"
                )
            if extra_in_new:
                error_parts.append(
                    f"Extra columns in new data: {sorted(extra_in_new)}"
                )

            raise ValueError(
                f"Schema mismatch for {existing_file}: {'; '.join(error_parts)}"
            )

        # Check data types match for critical columns
        type_mismatches = []
        for col_name in existing_columns:
            existing_type = existing_schema[col_name]
            new_type = new_schema[col_name]

            if existing_type != new_type:
                type_mismatches.append(f"{col_name}: {existing_type} -> {new_type}")

        if type_mismatches:
            raise ValueError(
                f"Data type mismatch for {existing_file}: {'; '.join(type_mismatches)}"
            )

    def _is_valid_month_format(self, month_str: str) -> bool:
        """Check if string is valid YYYY-MM month format."""
        if len(month_str) != 7 or month_str[4] != "-":
            return False
        try:
            year = int(month_str[:4])
            month = int(month_str[5:7])
            return 2020 <= year <= 2030 and 1 <= month <= 12
        except (ValueError, IndexError):
            return False


def create_parquet_data_service() -> ParquetDataService:
    """
    Factory function to create ParquetDataService.

    Returns:
        Configured ParquetDataService instance
    """
    return ParquetDataService()
