"""
Drop copier for time-shifting operations.

Copies and transforms a single drop (org/provider/month combination)
from source month to target month using manifest-based discovery.

Uses streaming S3 operations (no temp files) for better performance.
"""

import boto3
from typing import List
from ..models.time_shift_models import DropCopyResult
from .csv_date_transformer import CsvDateTransformer
from .manifest_manager import ManifestManager
from .billing_period_helper import BillingPeriodHelper


class DropCopier:
    """
    Copies and transforms a single drop.

    Uses manifest-based discovery to identify source drop and creates
    CloudZero-compatible target structure.
    """

    def __init__(
        self,
        csv_transformer: CsvDateTransformer,
        manifest_manager: ManifestManager,
        billing_helper: BillingPeriodHelper,
        s3_bucket: str,
        s3_root_folder: str = "piedpiper-custom-billing"
    ):
        """
        Initialize drop copier.

        Args:
            csv_transformer: CSV date transformer service
            manifest_manager: Manifest manager service
            billing_helper: Billing period helper service
            s3_bucket: S3 bucket name
            s3_root_folder: Root folder for billing data
        """
        self.csv_transformer = csv_transformer
        self.manifest_manager = manifest_manager
        self.billing_helper = billing_helper
        self.s3_bucket = s3_bucket
        self.s3_root_folder = s3_root_folder
        self.s3_client = boto3.client('s3')

    def copy_drop(
        self,
        org: str,
        provider: str,
        source_month: str,
        target_month: str
    ) -> DropCopyResult:
        """
        Copy and transform single drop using manifest-based discovery.

        Args:
            org: Organization name
            provider: Provider name
            source_month: Source month (YYYY-MM)
            target_month: Target month (YYYY-MM)

        Returns:
            DropCopyResult with statistics
        """
        import time
        start_time = time.time()

        try:
            # Convert months to billing periods
            source_period = self.billing_helper.month_to_billing_period(source_month)
            target_period = self.billing_helper.month_to_billing_period(target_month)

            # Read source manifest to get active drop ID
            source_drop_id = self._get_source_drop_id(org, provider, source_period)

            # Build source drop prefix from manifest
            source_prefix = self.billing_helper.build_drop_prefix(
                self.s3_root_folder, org, provider, source_period, source_drop_id
            )

            # List files in source drop
            csv_files = self._list_csv_files(source_prefix)
            if not csv_files:
                raise RuntimeError(f"No CSV files in {source_prefix}")

            # Generate target drop ID and prefix
            target_drop_id = self.billing_helper.generate_drop_id()
            target_prefix = self.billing_helper.build_drop_prefix(
                self.s3_root_folder, org, provider, target_period, target_drop_id
            )

            # Transform files
            total_records = self._transform_files(
                csv_files, source_prefix, target_prefix, source_month, target_month
            )

            # Create target manifest
            self._create_target_manifest(org, provider, target_period, target_drop_id)

            processing_time = time.time() - start_time

            return self._build_success_result(
                org, provider, source_month, target_month,
                source_period, target_period, source_drop_id, target_drop_id,
                len(csv_files), total_records, processing_time
            )

        except Exception as e:
            processing_time = time.time() - start_time
            return self._build_failure_result(
                org, provider, source_month, target_month, processing_time, str(e)
            )

    def _get_source_drop_id(
        self, org: str, provider: str, billing_period: str
    ) -> str:
        """
        Read source manifest to get active drop ID.

        Args:
            org: Organization name
            provider: Provider name
            billing_period: Billing period (YYYYMMDD-YYYYMMDD)

        Returns:
            Source drop ID from manifest

        Raises:
            RuntimeError: If manifest missing or invalid
        """
        manifest_key = self.billing_helper.build_manifest_key(
            self.s3_root_folder, org, provider, billing_period
        )
        manifest = self.manifest_manager.read_manifest(manifest_key)

        if 'current_drop_id' not in manifest:
            raise RuntimeError(
                f"Invalid manifest - missing current_drop_id: {manifest_key}"
            )

        return manifest['current_drop_id']

    def _list_csv_files(self, prefix: str) -> List[str]:
        """
        List all CSV.gz files in S3 prefix.

        Args:
            prefix: S3 prefix to search

        Returns:
            List of S3 keys for CSV files
        """
        csv_files = []
        paginator = self.s3_client.get_paginator('list_objects_v2')

        for page in paginator.paginate(Bucket=self.s3_bucket, Prefix=prefix):
            if 'Contents' not in page:
                continue

            for obj in page['Contents']:
                key = obj['Key']
                if key.endswith('.csv.gz') or key.endswith('.csv'):
                    csv_files.append(key)

        return csv_files

    def _transform_csv_file(
        self,
        source_key: str,
        target_key: str,
        source_month: str,
        target_month: str
    ):
        """
        Transform single CSV file via S3 streaming (no temp files).

        Streams directly from S3 → memory → S3.
        Faster and simpler than temp file approach.

        Args:
            source_key: Source S3 key
            target_key: Target S3 key
            source_month: Source month (YYYY-MM)
            target_month: Target month (YYYY-MM)

        Returns:
            TransformStats with record counts and timing
        """
        import time
        import polars as pl
        from io import BytesIO

        start_time = time.time()

        # Stream from S3 directly to DataFrame
        response = self.s3_client.get_object(
            Bucket=self.s3_bucket,
            Key=source_key
        )

        # Read compressed CSV from S3 stream
        # Force all columns to string to avoid schema inference issues
        # (e.g., cost center "CC-ENG-004" vs numeric IDs)
        df = pl.read_csv(
            response['Body'],
            has_header=True,
            infer_schema_length=0  # Read entire file as strings
        )

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

        # Write to buffer as compressed CSV
        buffer = BytesIO()
        df_transformed.write_csv(buffer)
        buffer.seek(0)

        # Compress buffer content
        import gzip
        compressed_buffer = BytesIO()
        with gzip.GzipFile(fileobj=compressed_buffer, mode='wb') as gz:
            gz.write(buffer.getvalue())
        compressed_buffer.seek(0)

        # Stream directly to S3
        self.s3_client.put_object(
            Bucket=self.s3_bucket,
            Key=target_key,
            Body=compressed_buffer.getvalue(),
            ContentType='text/csv',
            ContentEncoding='gzip'
        )

        # Update stats with total time (includes S3 I/O)
        transform_stats.processing_time_sec = time.time() - start_time

        return transform_stats

    def _transform_files(
        self,
        csv_files: List[str],
        source_prefix: str,
        target_prefix: str,
        source_month: str,
        target_month: str
    ) -> int:
        """
        Transform all CSV files from source to target.

        Args:
            csv_files: List of source CSV file S3 keys
            source_prefix: Source drop prefix
            target_prefix: Target drop prefix
            source_month: Source month (YYYY-MM)
            target_month: Target month (YYYY-MM)

        Returns:
            Total records processed
        """
        total_records = 0
        for csv_key in csv_files:
            # Extract filename and build target key
            filename = csv_key.split('/')[-1]
            target_key = f"{target_prefix}{filename}"

            stats = self._transform_csv_file(
                csv_key, target_key, source_month, target_month
            )
            total_records += stats.records_output

        return total_records

    def _create_target_manifest(
        self, org: str, provider: str, billing_period: str, drop_id: str
    ) -> None:
        """
        Create target manifest in CloudZero format.

        Args:
            org: Organization name
            provider: Provider name
            billing_period: Billing period (YYYYMMDD-YYYYMMDD)
            drop_id: Target drop ID
        """
        manifest_key = self.billing_helper.build_manifest_key(
            self.s3_root_folder, org, provider, billing_period
        )
        manifest_data = self.manifest_manager.create_manifest(drop_id)
        self.manifest_manager.write_manifest(manifest_key, manifest_data)

    def _build_success_result(
        self,
        org: str,
        provider: str,
        source_month: str,
        target_month: str,
        source_period: str,
        target_period: str,
        source_drop_id: str,
        target_drop_id: str,
        files_copied: int,
        total_records: int,
        processing_time: float
    ) -> DropCopyResult:
        """Build success result with all metadata."""
        return DropCopyResult(
            org=org,
            provider=provider,
            source_month=source_month,
            target_month=target_month,
            source_billing_period=source_period,
            target_billing_period=target_period,
            source_drop_id=source_drop_id,
            target_drop_id=target_drop_id,
            files_copied=files_copied,
            total_records=total_records,
            processing_time_sec=processing_time,
            success=True
        )

    def _build_failure_result(
        self,
        org: str,
        provider: str,
        source_month: str,
        target_month: str,
        processing_time: float,
        error_message: str
    ) -> DropCopyResult:
        """Build failure result with error."""
        return DropCopyResult(
            org=org,
            provider=provider,
            source_month=source_month,
            target_month=target_month,
            processing_time_sec=processing_time,
            success=False,
            error_message=error_message
        )
