from typing import List, Optional
"""
JsonlDataService - Handle JSONL operations for upload stage.

Primary Responsibility: JSONL file operations for CloudZero API upload

Capabilities:
- API payload generation (CloudZero API-compatible JSONL)
- Size-constrained chunking for API limits
- Async I/O for large datasets
- CloudZero format compliance
- Memory-efficient streaming writes
"""

import asyncio
from pathlib import Path

import polars as pl


class JsonlDataService:
    """
    Service for JSONL data operations for upload stage.

    Provides specialized JSONL operations for CloudZero API uploads:
    - API payload generation
    - Size-constrained chunking
    - CloudZero format compliance
    - Memory-efficient processing

    Features:
    - CloudZero API-compatible JSONL format
    - Automatic chunking for API size limits
    - Async I/O for non-blocking operations
    - Memory-efficient streaming for large datasets
    - Clear error reporting
    """

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

    def write_api_payload(self, data: pl.DataFrame, output_path: Path) -> None:
        """
        Write DataFrame as CloudZero API-compatible JSONL.

        Converts DataFrame to JSONL format suitable for CloudZero API upload.
        Removes internal columns and ensures proper formatting.

        Args:
            data: DataFrame to write as JSONL
            output_path: Target JSONL file path

        Raises:
            ValueError: If data is empty or invalid
        """
        if data is None or data.is_empty():
            # For empty data, create empty file
            output_path.parent.mkdir(parents=True, exist_ok=True)
            output_path.write_text("")
            return

        # Clean data for API upload
        cleaned_data = self._clean_data_for_api(data)

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

        try:
            # Write as JSONL
            jsonl_content = cleaned_data.write_ndjson()
            output_path.write_text(jsonl_content)
        except Exception as e:
            raise ValueError(f"Failed to write JSONL payload to {output_path}: {e}")

    async def write_api_payload_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 without blocking.

        Args:
            data: DataFrame to write
            output_path: Target JSONL file path
        """
        if data is None or data.is_empty():
            # Create parent directory and empty file
            output_path.parent.mkdir(parents=True, exist_ok=True)
            async with asyncio.to_thread(output_path.write_text, ""):
                pass
            return

        # Clean data for API upload
        cleaned_data = self._clean_data_for_api(data)

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

        # Convert to JSONL format
        jsonl_content = cleaned_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_content)

    def chunk_payload_by_size(
        self, data: pl.DataFrame, max_size_mb: float = 4.5
    ) -> List[pl.DataFrame]:
        """
        Split large datasets into size-constrained chunks.

        Splits DataFrames to respect CloudZero API size limits while
        maintaining memory efficiency.

        Args:
            data: DataFrame to split into chunks
            max_size_mb: Maximum size per chunk in MB

        Returns:
            List of DataFrame chunks
        """
        if data is None or data.is_empty():
            return []

        # Estimate row size for chunking
        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, data.height))

        chunks = []
        total_rows = data.height

        for start_idx in range(0, total_rows, target_rows_per_chunk):
            end_idx = min(start_idx + target_rows_per_chunk, total_rows)
            chunk = data.slice(start_idx, end_idx - start_idx)
            chunks.append(chunk)

        return chunks

    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)

        # Split data into chunks
        chunks = self.chunk_payload_by_size(data, max_size_mb)

        payload_files = []
        for chunk_idx, chunk_data in enumerate(chunks):
            # Create filename
            if provider and month:
                filename = f"{provider}_{month}_payload_{chunk_idx}.jsonl"
            else:
                filename = f"payload_{chunk_idx}.jsonl"

            payload_file = payloads_dir / filename

            # Write chunk as JSONL
            self.write_api_payload(chunk_data, payload_file)
            payload_files.append(payload_file)

        return payload_files

    def load_jsonl_data(self, file_path: Path) -> pl.DataFrame:
        """
        Load data from JSONL file.

        Args:
            file_path: Path to JSONL file

        Returns:
            Loaded DataFrame

        Raises:
            FileNotFoundError: If file doesn't exist
            ValueError: If file cannot be parsed
        """
        if not file_path.exists():
            raise FileNotFoundError(f"JSONL file not found: {file_path}")

        try:
            return pl.read_ndjson(file_path)
        except Exception as e:
            raise ValueError(f"Cannot parse JSONL file {file_path}: {e}")

    def validate_api_payload_format(self, data: pl.DataFrame) -> List[str]:
        """
        Validate DataFrame format for CloudZero API compatibility.

        Args:
            data: DataFrame to validate

        Returns:
            List of validation error messages (empty if valid)
        """
        errors = []

        if data is None:
            errors.append("Data cannot be None")
            return errors

        if data.is_empty():
            # Empty data is valid (results in empty payload)
            return errors

        # Check for required columns (basic CloudZero format)
        required_columns = ["cost"]  # Minimum requirement
        missing_columns = [col for col in required_columns if col not in data.columns]
        if missing_columns:
            errors.append(f"Missing required columns: {missing_columns}")

        # Check for internal columns that should be removed
        internal_columns = ["billing_month", "_merge_key"]
        present_internal = [col for col in internal_columns if col in data.columns]
        if present_internal:
            errors.append(f"Internal columns should be removed: {present_internal}")

        # Validate data types
        if "cost" in data.columns:
            cost_column = data["cost"]
            if not cost_column.dtype.is_numeric():
                errors.append("'cost' column must be numeric")

        return errors

    def _clean_data_for_api(self, data: pl.DataFrame) -> pl.DataFrame:
        """
        Clean DataFrame for CloudZero API upload.

        Removes internal columns and ensures proper formatting for API.

        Args:
            data: DataFrame to clean

        Returns:
            Cleaned DataFrame ready for API upload
        """
        if data.is_empty():
            return data

        cleaned_data = data

        # Remove internal columns if present
        internal_columns = ["billing_month", "_merge_key"]
        columns_to_remove = [
            col for col in internal_columns if col in cleaned_data.columns
        ]

        if columns_to_remove:
            cleaned_data = cleaned_data.drop(columns_to_remove)

        # Ensure cost column is present and numeric
        if "cost" in cleaned_data.columns:
            # Ensure cost is properly formatted
            cleaned_data = cleaned_data.with_columns(
                pl.col("cost").cast(pl.Float64, strict=False)
            )

        return cleaned_data

    def _estimate_jsonl_row_size(self, data: pl.DataFrame) -> float:
        """
        Estimate average row size in JSONL format.

        Args:
            data: DataFrame to estimate

        Returns:
            Estimated average row size in bytes
        """
        if data.is_empty():
            return 1000  # Default estimate

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

        # Clean sample data
        cleaned_sample = self._clean_data_for_api(sample_data)

        if cleaned_sample.is_empty():
            return 1000  # Default estimate

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

            # Add buffer for JSON overhead
            return avg_row_size * 1.2
        except Exception:
            return 1000  # Default estimate


def create_jsonl_data_service() -> JsonlDataService:
    """
    Factory function to create JsonlDataService.

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