"""
CSV date transformer using Polars vectorization.

CRITICAL: This service must follow strict Polars performance rules:
- NEVER use iter_rows() or to_list()
- Use vectorized operations only
- Use with_columns() for transformations
- Use filter() for filtering

Performance target: 10-100x faster than row-by-row processing.
"""

import polars as pl
from datetime import datetime
from typing import Tuple
from ..models.time_shift_models import TransformStats
from .day_count_calculator import DayCountCalculator


class CsvDateTransformer:
    """
    Transforms dates in CSV files using Polars vectorization.

    Handles:
    - Date parsing (vectorized)
    - Invalid day filtering (vectorized)
    - Date transformation (vectorized)
    - Compression (gzip)
    """

    def __init__(self, day_calculator: DayCountCalculator):
        """
        Initialize transformer with day calculator.

        Args:
            day_calculator: Calculator for valid days per month
        """
        self.day_calculator = day_calculator

    def transform_dataframe(
        self,
        df: pl.DataFrame,
        source_month: str,
        target_month: str
    ) -> Tuple[pl.DataFrame, TransformStats]:
        """
        Transform DataFrame with date shifting (in-memory, no files).

        CRITICAL: All operations are vectorized for performance.
        Use this for streaming S3 operations (no temp files).

        Args:
            df: Input DataFrame with time/usage_start column
            source_month: Source month (YYYY-MM)
            target_month: Target month (YYYY-MM)

        Returns:
            Tuple of (transformed_df, stats)

        Examples:
            >>> transformer = CsvDateTransformer(DayCountCalculator())
            >>> df = pl.read_csv("data.csv.gz")
            >>> df_transformed, stats = transformer.transform_dataframe(
            ...     df, "2025-01", "2026-01"
            ... )
        """
        import time
        start_time = time.time()

        records_input = len(df)

        # Parse date column (vectorized)
        df = self._parse_date_column(df)

        # Filter records for target month (vectorized)
        df = self._filter_by_target_month(df, target_month)

        # Transform dates (vectorized)
        df = self._shift_dates(df, source_month, target_month)

        processing_time = time.time() - start_time
        records_output = len(df)

        stats = TransformStats(
            records_input=records_input,
            records_output=records_output,
            processing_time_sec=processing_time
        )

        return df, stats

    def transform_file(
        self,
        source_path: str,
        target_path: str,
        source_month: str,
        target_month: str
    ) -> TransformStats:
        """
        Transform single CSV file with date shifting (file-based).

        Legacy method for local CLI execution.
        For ECS/API workers, use transform_dataframe() with streaming.

        Args:
            source_path: Path to source CSV.gz file
            target_path: Path to target CSV.gz file
            source_month: Source month (YYYY-MM)
            target_month: Target month (YYYY-MM)

        Returns:
            TransformStats with transformation metrics

        Examples:
            >>> transformer = CsvDateTransformer(DayCountCalculator())
            >>> stats = transformer.transform_file(
            ...     "s3://bucket/2025-01/data.csv.gz",
            ...     "s3://bucket/2026-01/data.csv.gz",
            ...     "2025-01",
            ...     "2026-01"
            ... )
        """
        # Read CSV.gz (Polars handles compression automatically)
        df = pl.read_csv(source_path)

        # Transform in memory
        df_transformed, stats = self.transform_dataframe(
            df, source_month, target_month
        )

        # Write CSV.gz (write uncompressed, then compress)
        import gzip
        import os
        temp_csv = target_path.replace('.gz', '') if target_path.endswith('.gz') else target_path + '.tmp'
        df_transformed.write_csv(temp_csv)

        # Compress if target is .gz
        if target_path.endswith('.gz'):
            with open(temp_csv, 'rb') as f_in:
                with gzip.open(target_path, 'wb') as f_out:
                    f_out.writelines(f_in)
            os.remove(temp_csv)

        return stats

    def _parse_date_column(self, df: pl.DataFrame) -> pl.DataFrame:
        """
        Parse date column into datetime type (vectorized).

        Args:
            df: Input DataFrame

        Returns:
            DataFrame with parsed_date column added
        """
        # Vectorized date parsing - processes all rows at once
        return df.with_columns([
            pl.col("time/usage_start")
            .str.strptime(pl.Datetime, "%Y-%m-%d %H:%M:%S%.f")
            .alias("parsed_date")
        ])

    def _filter_by_target_month(
        self,
        df: pl.DataFrame,
        target_month: str
    ) -> pl.DataFrame:
        """
        Filter records to valid days for target month (vectorized).

        Example: When transforming to February (28/29 days),
        filter out records from day 30 and 31.

        Args:
            df: DataFrame with parsed_date column
            target_month: Target month (YYYY-MM)

        Returns:
            Filtered DataFrame
        """
        year, month = self._parse_month(target_month)
        max_days = self.day_calculator.get_valid_days(year, month)

        # Vectorized filtering - no loops!
        # Extract day from all dates at once, then filter
        return df.filter(pl.col("parsed_date").dt.day() <= max_days)

    def _shift_dates(
        self,
        df: pl.DataFrame,
        source_month: str,
        target_month: str
    ) -> pl.DataFrame:
        """
        Shift all dates by calculated offset (vectorized).

        Args:
            df: DataFrame with parsed_date column
            source_month: Source month (YYYY-MM)
            target_month: Target month (YYYY-MM)

        Returns:
            DataFrame with transformed date column
        """
        # Calculate month offset
        offset_months = self._calculate_month_offset(source_month, target_month)

        # Vectorized date transformation
        # Apply offset to all dates at once, then format
        return df.with_columns([
            (
                pl.col("parsed_date")
                .dt.offset_by(f"{offset_months}mo")
                .dt.strftime("%Y-%m-%d %H:%M:%S.%6f")
                .alias("time/usage_start")
            )
        ]).drop("parsed_date")

    def _parse_month(self, month: str) -> Tuple[int, int]:
        """
        Parse month string to year and month integers.

        Args:
            month: Month in YYYY-MM format

        Returns:
            Tuple of (year, month)
        """
        parts = month.split("-")
        return int(parts[0]), int(parts[1])

    def _calculate_month_offset(
        self,
        source_month: str,
        target_month: str
    ) -> int:
        """
        Calculate month offset between source and target.

        Args:
            source_month: Source month (YYYY-MM)
            target_month: Target month (YYYY-MM)

        Returns:
            Number of months difference (can be negative)
        """
        source_year, source_mon = self._parse_month(source_month)
        target_year, target_mon = self._parse_month(target_month)

        # Calculate total months from reference point
        source_total = source_year * 12 + source_mon
        target_total = target_year * 12 + target_mon

        return target_total - source_total
