"""GCP-specific discount processors."""

import polars as pl

from .base import DiscountProcessor


class GCPCUDProcessor(DiscountProcessor):
    """GCP Committed Use Discount (CUD) processor."""

    def __init__(self):
        super().__init__(
            discount_percentage=0.30, discount_range=(0.30, 0.70)  # 30% of eligible resources  # 30-70% discount
        )

    def get_relevant_services(self) -> list[str]:
        return [
            "Compute Engine",
            "Cloud SQL",
            "Cloud Run",
            "BigQuery",
            "Cloud Spanner",
            "Cloud Dataflow",
        ]

    def _apply_specific_discount(self, data: pl.DataFrame, discount_mask: pl.Series) -> pl.DataFrame:
        """Apply GCP CUD specific discount logic."""
        result = data.clone()

        # Check if any rows already have 'Commitment' in usage_family
        if "usage_family" in result.columns:
            commitment_mask = result.select(pl.col("usage_family").str.contains("Commitment")).to_series()
            has_commitment_records = commitment_mask.any()

            if has_commitment_records:
                # Update the discount_mask to only include rows that have 'Commitment' in usage_family
                discount_mask = discount_mask & commitment_mask
            else:
                # No existing commitment records, apply the new commitment value
                result = result.with_columns(
                    pl.when(pl.Series(discount_mask))
                    .then(pl.lit("Commitment v1: Cpu in Americas for 1 Year"))
                    .otherwise(pl.col("usage_family"))
                    .alias("usage_family")
                )
        else:
            # usage_family column doesn't exist, create it with the new commitment value
            result = result.with_columns(
                pl.when(pl.Series(discount_mask))
                .then(pl.lit("Commitment v1: Cpu in Americas for 1 Year"))
                .otherwise(pl.lit(None))
                .alias("usage_family")
            )

        # Apply discount calculation
        mask_count = discount_mask.sum()
        if mask_count == 0:
            # Add default discount columns
            result = result.with_columns([pl.lit(0.0).alias("discount_amount"), pl.lit(None).alias("discount_type")])
            return result

        discount_factors = self._generate_discount_factor(mask_count)
        if mask_count == 1:
            discount_factors = [discount_factors]

        # Create factor mapping for polars operations
        factor_index = 0
        factor_mapping = {}
        for i, should_discount in enumerate(discount_mask):
            if should_discount:
                factor_mapping[i] = discount_factors[factor_index]
                factor_index += 1

        # Apply CUD discount logic
        result = result.with_columns(
            [
                # Calculate discount amount (cost reduction)
                pl.when(pl.Series(discount_mask))
                .then(pl.col("cost") * pl.Series([factor_mapping.get(i, 0.0) for i in range(len(discount_mask))]))
                .otherwise(0.0)
                .alias("discount_amount"),
                # Set discount type
                pl.when(pl.Series(discount_mask)).then(pl.lit("GCP CUD")).otherwise(None).alias("discount_type"),
            ]
        )

        return result


class GCPSpotProcessor(DiscountProcessor):
    """GCP Spot VM (Preemptible) discount processor."""

    def __init__(self):
        super().__init__(
            discount_percentage=0.15,  # 15% of eligible resources
            discount_range=(0.1, 0.3),  # 10-30% discount (70-90% savings)
        )

    def get_relevant_services(self) -> list[str]:
        return ["Compute Engine", "Cloud Dataflow", "Dataproc"]

    def _apply_specific_discount(self, data: pl.DataFrame, discount_mask: pl.Series) -> pl.DataFrame:
        """Apply GCP Spot VM specific discount logic."""
        result = data.clone()

        # Set usage_family to Spot for selected rows
        result = result.with_columns(
            pl.when(pl.Series(discount_mask))
            .then(pl.lit("Spot"))
            .otherwise(pl.col("usage_family") if "usage_family" in result.columns else pl.lit(None))
            .alias("usage_family")
        )

        # Apply discount calculation
        mask_count = discount_mask.sum()
        if mask_count == 0:
            # Add default discount columns
            result = result.with_columns([pl.lit(0.0).alias("discount_amount"), pl.lit(None).alias("discount_type")])
            return result

        discount_factors = self._generate_discount_factor(mask_count)
        if mask_count == 1:
            discount_factors = [discount_factors]

        # Create factor mapping for polars operations
        factor_index = 0
        factor_mapping = {}
        for i, should_discount in enumerate(discount_mask):
            if should_discount:
                factor_mapping[i] = discount_factors[factor_index]
                factor_index += 1

        # Apply Spot discount logic
        result = result.with_columns(
            [
                # Calculate discount amount (cost reduction)
                pl.when(pl.Series(discount_mask))
                .then(pl.col("cost") * pl.Series([factor_mapping.get(i, 0.0) for i in range(len(discount_mask))]))
                .otherwise(0.0)
                .alias("discount_amount"),
                # Set discount type
                pl.when(pl.Series(discount_mask)).then(pl.lit("GCP Spot")).otherwise(None).alias("discount_type"),
            ]
        )

        return result
