from typing import (
    Dict,
    List,
    Optional,
    Tuple,
    Union,
)
"""
Data Service - Pure data loading and saving operations.

Primary Responsibility: Load and save data files with format detection

Focus Areas:
- Billing data: Parquet files, CSV files containing cost/usage data
- Reference data: Provider configuration data (providers.csv)
- Stage output data: Generated synthetic billing data, processed pipeline data
- Payload data: JSONL files for CloudZero API upload

Does NOT handle: Schema alignment, processing strategy selection, performance monitoring

This service preserves all existing data I/O behavior including:
- Multi-format support (Parquet, CSV, JSONL)
- Compression options (snappy, lz4, none)
- Partitioned data loading and saving
- Provider/month directory structure handling
- Atomic file operations for data safety
- Format detection and fallback mechanisms
"""

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

import polars as pl


class DataFormat(Enum):
    """Supported data formats."""

    PARQUET = "parquet"
    CSV = "csv"
    JSONL = "jsonl"


class CompressionType(Enum):
    """Supported compression types."""

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


@dataclass
class DataFileInfo:
    """
    Information about a data file.

    Contains metadata about data files including format, size,
    compression, and schema information.
    """

    file_path: Path
    format: DataFormat
    compression: Optional[CompressionType]
    size_bytes: int
    row_count: Optional[int] = None
    column_count: Optional[int] = None
    schema: Optional[Dict[str, str]] = None

    @property
    def size_mb(self) -> float:
        """Get file size in megabytes."""
        return self.size_bytes / (1024 * 1024)


@dataclass
class SaveOptions:
    """
    Options for saving data files.

    Configures compression, format, and other save parameters.
    """

    format: DataFormat = DataFormat.PARQUET
    compression: CompressionType = CompressionType.SNAPPY
    create_directories: bool = True
    atomic_write: bool = True


class DataService:
    """
    Service for pure data loading and saving operations.

    Provides centralized data I/O operations while preserving all existing
    file handling behavior, format support, and error handling patterns.

    Key Features:
    - Multi-format support (Parquet, CSV, JSONL)
    - Format detection with fallback mechanisms
    - Compression handling (snappy, lz4, none)
    - Partitioned data loading and saving
    - Provider/month directory structure support
    - Atomic file operations for data safety
    - Schema validation and error reporting
    """

    def __init__(self):
        """Initialize DataService."""
        # Cache for schema information to avoid repeated scanning
        self._schema_cache: Dict[str, Dict[str, str]] = {}

    # === Core Data Loading Operations ===

    def load_data(self, source: Path) -> pl.DataFrame:
        """
        Load data from a file with automatic format detection.

        Supports all standard data formats with automatic fallback
        behavior. Preserves existing loading patterns and error handling.

        Args:
            source: Path to data file or directory

        Returns:
            Loaded DataFrame

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

        if source.is_dir():
            return self._load_directory(source)
        else:
            return self._load_single_file(source)

    def load_partitioned_data(
        self, base_dir: Path, partition_pattern: str = "provider"
    ) -> pl.DataFrame:
        """
        Load partitioned data with schema alignment.

        Handles different partitioning patterns:
        - Provider partitioning: AWS.parquet, GCP.parquet, Azure.parquet
        - Month partitioning: Provider/YYYY-MM/data.parquet

        Args:
            base_dir: Base directory containing partitioned data
            partition_pattern: "provider", "month", or "auto"

        Returns:
            Combined DataFrame with aligned schema
        """
        if not base_dir.exists():
            raise FileNotFoundError(f"Partition directory not found: {base_dir}")

        if partition_pattern == "auto":
            partition_pattern = self._detect_partition_pattern(base_dir)

        if partition_pattern == "provider":
            return self._load_provider_partitioned_data(base_dir)
        elif partition_pattern == "month":
            return self._load_month_partitioned_data(base_dir)
        else:
            raise ValueError(f"Unknown partition pattern: {partition_pattern}")

    def detect_format(self, file_path: Path) -> DataFormat:
        """
        Detect file format from path and content.

        Uses file extension for initial detection with content-based
        fallback for ambiguous cases.

        Args:
            file_path: Path to file

        Returns:
            Detected data format
        """
        if not file_path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        # Primary detection by extension
        extension = file_path.suffix.lower()
        if extension == ".parquet":
            return DataFormat.PARQUET
        elif extension == ".csv":
            return DataFormat.CSV
        elif extension in [".jsonl", ".ndjson"]:
            return DataFormat.JSONL
        else:
            # Fallback: try CSV for unknown extensions
            return DataFormat.CSV

    def get_file_info(self, file_path: Path) -> DataFileInfo:
        """
        Get comprehensive information about a data file.

        Analyzes file metadata including format, size, and schema
        without loading the entire file into memory.

        Args:
            file_path: Path to data file

        Returns:
            DataFileInfo with file metadata
        """
        if not file_path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        format_type = self.detect_format(file_path)
        size_bytes = file_path.stat().st_size

        # Get row and column counts efficiently
        row_count, column_count, schema = self._get_file_stats(file_path, format_type)

        # Detect compression from file metadata
        compression = self._detect_compression(file_path, format_type)

        return DataFileInfo(
            file_path=file_path,
            format=format_type,
            compression=compression,
            size_bytes=size_bytes,
            row_count=row_count,
            column_count=column_count,
            schema=schema,
        )

    # === Core Data Saving Operations ===

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

        Handles compression, format conversion, and atomic writes.
        Preserves existing saving patterns and safety mechanisms.

        Args:
            data: DataFrame to save
            destination: Target file path
            options: Save options (format, 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 = SaveOptions()

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

        # Use atomic write pattern for safety
        if options.atomic_write:
            self._atomic_save(data, destination, options)
        else:
            self._direct_save(data, destination, options)

    def save_partitioned_data(
        self,
        data: pl.DataFrame,
        base_dir: Path,
        partition_column: str,
        options: Optional[SaveOptions] = None,
    ) -> Dict[str, Path]:
        """
        Save data partitioned by specified column.

        Creates separate files for each unique value in the partition column.
        Commonly used for provider-based partitioning.

        Args:
            data: DataFrame to partition and save
            base_dir: Base directory for partitioned files
            partition_column: Column to partition by (e.g., "cloud_provider")
            options: Save options

        Returns:
            Dictionary mapping partition values to file paths
        """
        if partition_column not in data.columns:
            raise ValueError(f"Partition column '{partition_column}' not found in data")

        if options is None:
            options = SaveOptions()

        # Create base directory
        if options.create_directories:
            base_dir.mkdir(parents=True, exist_ok=True)

        partition_files = {}
        partition_values = data[partition_column].unique().to_list()

        for value in partition_values:
            if value is None:
                continue

            # Filter data for this partition
            partition_data = data.filter(pl.col(partition_column) == value)

            # Create partition file path
            safe_value = str(value).replace("/", "_").replace(" ", "_")
            partition_file = base_dir / f"{safe_value}.{options.format.value}"

            # Save partition
            self.save_data(partition_data, partition_file, options)
            partition_files[str(value)] = partition_file

        return partition_files

    def save_provider_month_data(
        self,
        data: pl.DataFrame,
        base_dir: Path,
        provider: str,
        month: str,
        options: Optional[SaveOptions] = None,
    ) -> Path:
        """
        Save data in provider/month directory structure.

        Creates the standard provider/month organization:
        {base_dir}/{provider}/{month}/data.{format}

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

        Returns:
            Path to saved file
        """
        if options is None:
            options = SaveOptions()

        # 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)

        # Save data file
        data_file = provider_month_dir / f"data.{options.format.value}"
        self.save_data(data, data_file, options)

        return data_file

    # === JSONL Operations (for API uploads) ===

    async def write_jsonl_async(self, data: pl.DataFrame, output_path: Path) -> None:
        """
        Write DataFrame as JSONL using async I/O.

        Optimized for large datasets with non-blocking I/O operations.
        Used for generating API upload payloads.

        Args:
            data: DataFrame to write
            output_path: Target JSONL file path
        """
        if data.is_empty():
            return

        # Create parent directory if needed
        output_path.parent.mkdir(parents=True, exist_ok=True)

        # Convert to JSONL format
        jsonl_data = data.write_ndjson()

        # Write using async I/O to avoid blocking
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, output_path.write_text, jsonl_data)

    def write_payload_chunks(
        self,
        data: pl.DataFrame,
        output_dir: Path,
        max_size_mb: float = 4.5,
        provider: Optional[str] = None,
        month: Optional[str] = None,
    ) -> List[Path]:
        """
        Write DataFrame as chunked JSONL payload files.

        Splits large datasets into multiple files respecting CloudZero API
        size limits while maintaining memory efficiency.

        Args:
            data: DataFrame to write as payload chunks
            output_dir: Directory to write payload files
            max_size_mb: Maximum size per payload file in MB
            provider: Provider name for filename organization (optional)
            month: Month for filename organization (optional)

        Returns:
            List of paths to created payload files
        """
        if data is None or data.is_empty():
            return []

        # Create payloads directory
        payloads_dir = output_dir / "payloads"
        payloads_dir.mkdir(parents=True, exist_ok=True)

        # Calculate chunk size based on target file size
        estimated_row_size = self._estimate_jsonl_row_size(data)
        target_rows_per_chunk = int((max_size_mb * 1024 * 1024) / estimated_row_size)
        target_rows_per_chunk = max(1, min(target_rows_per_chunk, len(data)))

        payload_files = []
        total_rows = len(data)

        for chunk_idx in range(0, total_rows, target_rows_per_chunk):
            chunk_end = min(chunk_idx + target_rows_per_chunk, total_rows)
            chunk_data = data.slice(chunk_idx, chunk_end - chunk_idx)

            # Create filename
            if provider and month:
                filename = f"{provider}_{month}_payload_{len(payload_files)}.jsonl"
            else:
                filename = f"payload_{len(payload_files)}.jsonl"

            payload_file = payloads_dir / filename

            # Remove internal columns if present
            if "billing_month" in chunk_data.columns:
                chunk_data = chunk_data.drop("billing_month")
            if "cloud_provider" in chunk_data.columns and provider:
                # Keep cloud_provider if it's the partitioning column
                pass

            # Write payload file
            asyncio.run(self.write_jsonl_async(chunk_data, payload_file))
            payload_files.append(payload_file)

        return payload_files

    # === Schema Operations ===

    def validate_schema(self, data: pl.DataFrame, required_columns: List[str]) -> None:
        """
        Validate DataFrame schema against requirements.

        Args:
            data: DataFrame to validate
            required_columns: List of required column names

        Raises:
            ValueError: If schema validation fails
        """
        missing_columns = [col for col in required_columns if col not in data.columns]
        if missing_columns:
            raise ValueError(
                f"Missing required columns: {missing_columns}. "
                f"Available columns: {data.columns}"
            )

    def align_schemas(self, dataframes: List[pl.DataFrame]) -> List[pl.DataFrame]:
        """
        Align schemas across multiple DataFrames.

        Ensures all DataFrames have the same columns in the same order,
        adding missing columns as null values with appropriate types.

        Args:
            dataframes: List of DataFrames to align

        Returns:
            List of DataFrames with aligned schemas
        """
        if not dataframes:
            return []

        if len(dataframes) == 1:
            return dataframes

        # Collect all unique columns across all DataFrames
        all_columns = set()
        schemas = {}

        for df in dataframes:
            all_columns.update(df.columns)
            for col_name, col_type in df.schema.items():
                if col_name not in schemas:
                    schemas[col_name] = col_type

        all_columns = sorted(all_columns)

        # Align each DataFrame
        aligned_dataframes = []
        for df in dataframes:
            aligned_df = df

            # Add missing columns
            for col in all_columns:
                if col not in aligned_df.columns:
                    null_type = schemas.get(col, pl.Utf8)
                    aligned_df = aligned_df.with_columns(
                        pl.lit(None, dtype=null_type).alias(col)
                    )

            # Ensure consistent column order
            aligned_df = aligned_df.select(all_columns)
            aligned_dataframes.append(aligned_df)

        return aligned_dataframes

    # === Private Helper Methods ===

    def _load_single_file(self, file_path: Path) -> pl.DataFrame:
        """Load data from a single file with format detection."""
        format_type = self.detect_format(file_path)

        if format_type == DataFormat.PARQUET:
            return pl.read_parquet(file_path)
        elif format_type == DataFormat.CSV:
            return pl.read_csv(file_path)
        elif format_type == DataFormat.JSONL:
            return pl.read_ndjson(file_path)
        else:
            # Fallback to CSV for unknown formats
            if format_type != DataFormat.CSV:
                try:
                    return pl.read_csv(file_path)
                except Exception:
                    raise ValueError(f"Unsupported format: {format_type}")
            raise ValueError(f"Unsupported format: {format_type}")

    def _load_directory(self, directory: Path) -> pl.DataFrame:
        """Load all compatible files from a directory."""
        parquet_files = list(directory.glob("**/*.parquet"))

        if parquet_files:
            return pl.read_parquet(parquet_files)

        csv_files = list(directory.glob("**/*.csv"))
        if csv_files:
            dataframes = [self._load_single_file(f) for f in csv_files]
            if len(dataframes) == 1:
                return dataframes[0]
            else:
                aligned_dfs = self.align_schemas(dataframes)
                return pl.concat(aligned_dfs, how="vertical")

        raise FileNotFoundError(f"No compatible data files found in {directory}")

    def _load_provider_partitioned_data(self, base_dir: Path) -> pl.DataFrame:
        """Load provider-partitioned data (AWS.parquet, GCP.parquet, etc.)."""
        provider_files = []
        for provider in ["AWS", "Azure", "GCP"]:
            provider_file = base_dir / f"{provider}.parquet"
            if provider_file.exists():
                provider_files.append(provider_file)

        if not provider_files:
            raise FileNotFoundError(f"No provider partition files found in {base_dir}")

        dataframes = [self._load_single_file(f) for f in provider_files]
        aligned_dfs = self.align_schemas(dataframes)
        return pl.concat(aligned_dfs, how="vertical")

    def _load_month_partitioned_data(self, base_dir: Path) -> pl.DataFrame:
        """Load month-partitioned data (Provider/YYYY-MM/data.parquet)."""
        data_files = []

        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():
                            data_files.append(data_file)

        if not data_files:
            raise FileNotFoundError(f"No month partition files found in {base_dir}")

        dataframes = [self._load_single_file(f) for f in data_files]
        aligned_dfs = self.align_schemas(dataframes)
        return pl.concat(aligned_dfs, how="vertical")

    def _detect_partition_pattern(self, base_dir: Path) -> str:
        """Detect partitioning pattern in directory."""
        # Check for provider files (AWS.parquet, etc.)
        provider_files = any(
            (base_dir / f"{provider}.parquet").exists()
            for provider in ["AWS", "Azure", "GCP"]
        )

        if provider_files:
            return "provider"

        # Check for month structure (Provider/YYYY-MM/)
        has_month_structure = False
        for item in base_dir.iterdir():
            if item.is_dir() and item.name in ["AWS", "Azure", "GCP"]:
                for subitem in item.iterdir():
                    if subitem.is_dir() and self._is_valid_month_format(subitem.name):
                        has_month_structure = True
                        break
                if has_month_structure:
                    break

        if has_month_structure:
            return "month"

        return "provider"  # Default fallback

    def _atomic_save(
        self, data: pl.DataFrame, destination: Path, options: SaveOptions
    ) -> None:
        """Save data using atomic write pattern (temp file + rename)."""
        temp_file = destination.with_suffix(f".tmp{destination.suffix}")

        try:
            self._direct_save(data, temp_file, options)
            temp_file.replace(destination)
        finally:
            # Clean up temp file if it still exists
            if temp_file.exists():
                temp_file.unlink(missing_ok=True)

    def _direct_save(
        self, data: pl.DataFrame, destination: Path, options: SaveOptions
    ) -> None:
        """Save data directly to destination."""
        compression = (
            options.compression.value
            if options.compression != CompressionType.NONE
            else None
        )

        if options.format == DataFormat.PARQUET:
            data.write_parquet(destination, compression=compression)
        elif options.format == DataFormat.CSV:
            data.write_csv(destination)
        elif options.format == DataFormat.JSONL:
            jsonl_data = data.write_ndjson()
            destination.write_text(jsonl_data)
        else:
            raise ValueError(f"Unsupported save format: {options.format}")

    def _get_file_stats(
        self, file_path: Path, format_type: DataFormat
    ) -> Optional[Tuple[Union[int, None], Union[int, None], Dict[str, str]]]:
        """Get file statistics efficiently without loading entire file."""
        if format_type == DataFormat.PARQUET:
            # Use lazy frame to get metadata
            lazy_df = pl.scan_parquet(file_path)
            schema = dict(lazy_df.schema)

            # Get row count efficiently
            row_count = lazy_df.select(pl.len()).collect().item()

            return row_count, len(schema), schema

        elif format_type == DataFormat.CSV:
            # For CSV, we need to sample to get schema
            sample_df = pl.read_csv(file_path, n_rows=1)
            schema = dict(sample_df.schema)
            return None, len(schema), schema  # Row count expensive for CSV

        else:
            return None, None, None

    def _detect_compression(
        self, file_path: Path, format_type: DataFormat
    ) -> Optional[CompressionType]:
        """Detect compression type from file metadata."""
        if format_type == DataFormat.PARQUET:
            # Could inspect parquet metadata, but for now return None
            return None
        else:
            return None

    def _is_valid_month_format(self, month_str: str) -> bool:
        """Check if string is valid YYYY-MM month format."""
        try:
            if len(month_str) != 7 or month_str[4] != "-":
                return False
            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 _estimate_jsonl_row_size(self, data: pl.DataFrame) -> float:
        """Estimate average row size in JSONL format."""
        if data.is_empty():
            return 1000  # Default estimate

        # Sample first few rows to estimate size
        sample_size = min(10, len(data))
        sample_data = data.head(sample_size)

        # Convert sample to JSONL and measure
        sample_jsonl = sample_data.write_ndjson()
        avg_row_size = len(sample_jsonl.encode("utf-8")) / sample_size

        # Add buffer for JSON overhead
        return avg_row_size * 1.2
