"""Azure-specific discount processors."""

from typing import List
import numpy as np
import polars as pl

from .base import DiscountProcessor


class AzureSpotProcessor(DiscountProcessor):
    """Azure Spot VM 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 ["microsoft.compute", "microsoft.containerservice"]

    def get_additional_filters(self, data: pl.DataFrame) -> pl.Series:
        """Azure Spot VMs require pricing_unit to be '1 Hour'."""
        if "pricing_unit" not in data.columns:
            self.logger.warning("pricing_unit column missing for Azure Spot filtering, skipping pricing unit filter")
            return pl.Series([True] * len(data))
        
        return data.select(
            pl.col("pricing_unit") == "1 Hour"
        ).to_series()

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

        # Generate discount factors for selected rows
        mask_count = discount_mask.sum()
        if mask_count == 0:
            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 discount amount calculation
        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("Azure Spot"))
            .otherwise(None)
            .alias("discount_type"),
            
            # Update Azure-specific metadata
            pl.when(pl.Series(discount_mask))
            .then(pl.lit("Spot"))
            .otherwise(pl.col("operation"))
            .alias("operation")
        ])

        return result


class AzureReservationProcessor(DiscountProcessor):
    """Azure Reservation discount processor (currently disabled)."""

    def __init__(self):
        super().__init__(
            discount_percentage=0.25,  # 25% of eligible resources
            discount_range=(0.20, 0.45)  # 20-45% discount
        )

    def get_relevant_services(self) -> List[str]:
        return ["microsoft.compute", "microsoft.sql", "microsoft.storage"]

    def get_additional_filters(self, data: pl.DataFrame) -> pl.Series:
        """Azure Reservations require pricing_unit to be '1 Hour'."""
        if "pricing_unit" not in data.columns:
            self.logger.warning("pricing_unit column missing for Azure Reservation filtering, skipping pricing unit filter")
            return pl.Series([True] * len(data))
        
        return data.select(
            pl.col("pricing_unit") == "1 Hour"
        ).to_series()

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

        # Generate discount factors
        mask_count = discount_mask.sum()
        if mask_count == 0:
            return result

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

        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 reservation discount logic
        result = result.with_columns([
            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"),
            
            pl.when(pl.Series(discount_mask))
            .then(pl.lit("Azure Reservation"))
            .otherwise(None)
            .alias("discount_type"),
            
            # Reservation-specific metadata updates
            pl.when(pl.Series(discount_mask))
            .then(pl.lit("DiscountedUsage"))
            .otherwise(pl.col("line_item_type"))
            .alias("line_item_type"),
            
            pl.when(pl.Series(discount_mask))
            .then(pl.lit("Reservation"))
            .otherwise(pl.col("operation"))
            .alias("operation")
        ])

        return result


class AzureSavingsPlanProcessor(DiscountProcessor):
    """Azure Savings Plan discount processor (currently disabled)."""

    def __init__(self):
        super().__init__(
            discount_percentage=0.30,  # 30% of eligible resources
            discount_range=(0.15, 0.35)  # 15-35% discount
        )

    def get_relevant_services(self) -> List[str]:
        return ["microsoft.compute", "microsoft.containerservice"]

    def get_additional_filters(self, data: pl.DataFrame) -> pl.Series:
        """Azure Savings Plans require pricing_unit to be '1 Hour'."""
        if "pricing_unit" not in data.columns:
            self.logger.warning("pricing_unit column missing for Azure Savings Plan filtering, skipping pricing unit filter")
            return pl.Series([True] * len(data))
        
        return data.select(
            pl.col("pricing_unit") == "1 Hour"
        ).to_series()

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

        # Generate discount factors
        mask_count = discount_mask.sum()
        if mask_count == 0:
            return result

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

        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 Savings Plan discount logic
        result = result.with_columns([
            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"),
            
            pl.when(pl.Series(discount_mask))
            .then(pl.lit("Azure Savings Plan"))
            .otherwise(None)
            .alias("discount_type"),
            
            # Savings Plan metadata updates
            pl.when(pl.Series(discount_mask))
            .then(pl.lit("SavingsPlanCoveredUsage"))
            .otherwise(pl.col("line_item_type"))
            .alias("line_item_type")
        ])

        return result