"""Multi-pass telemetry data generator.

Four-pass architecture:
1. Even distribution - All elements start equal
2. Dominant element - One element gets 60% share (or varied)
3. Variance - Apply ±variance_pct noise
4. Cost isolation & anomalies - Calculate target shares, apply spikes
"""

from datetime import datetime, timedelta
from typing import List, Set, Dict

import polars as pl

from src.shared.anomaly import (
    calculate_anomaly_timing,
    calculate_anomaly_occurrences_in_job,
)
from .anomalies import (
    apply_daily_spike,
    apply_daily_step_up,
    apply_daily_step_down,
    apply_hourly_spike,
    apply_hourly_step_up,
    apply_hourly_step_down,
)
from .config import TelemetryConfig
from .cost_isolation import CostIsolationService
from .distribution import EvenDistributionService
from .dominance import DominantElementService
from .variance import TelemetryVarianceService

# Maximum telemetry increase for anomaly element (prevents unrealistic spikes)
# e.g., 10.0 = 10x increase = +900% (realistic for business activity)
MAX_TELEMETRY_MULTIPLIER = 10.0

# Anomaly type to function mapping
ANOMALY_FUNCTIONS = {
    "daily_spike": apply_daily_spike,
    "hourly_spike": apply_hourly_spike,
    "daily_step_up": apply_daily_step_up,
    "daily_step_down": apply_daily_step_down,
    "hourly_step_up": apply_hourly_step_up,
    "hourly_step_down": apply_hourly_step_down,
}


class TelemetryGenerator:
    """Orchestrates multi-pass telemetry generation.

    Single responsibility: Coordinate four-pass generation process.
    Delegates actual work to focused services.
    """

    def __init__(self, config: TelemetryConfig):
        """Initialize generator with config and services.

        Args:
            config: Telemetry configuration
        """
        self.config = config
        self.elements = list(config.elements.keys())
        self.base_seed = 42

        # Initialize services (dependency injection)
        self.distribution_service = EvenDistributionService()
        self.dominance_service = DominantElementService(
            dominant_element_id=config.get_dominant_element(),
            seed=self.base_seed
        )
        self.variance_service = TelemetryVarianceService(
            default_variance_pct=config.variance_pct
        )
        self.cost_isolation_service = CostIsolationService()

    def generate(self) -> pl.DataFrame:
        """Generate telemetry DataFrame using four-pass architecture.

        Returns:
            DataFrame with columns: timestamp, element_id, unit_value
        """
        # Create hourly timestamps
        timestamps = self._create_timestamps()

        # Pass 1: Even distribution
        df = self.distribution_service.generate_base_telemetry(
            timestamps, self.elements, self.config.typical_hourly_volume
        )

        # Pass 2: Dominant element redistribution
        df = self.dominance_service.apply_dominance(df)

        # Pass 3: Variance application (per-element)
        df = self._apply_variance_to_elements(df)

        # Pass 4: Anomaly application (per-element)
        df = self._apply_anomalies_to_elements(df)

        return df

    def _create_timestamps(self) -> List[datetime]:
        """Create list of hourly timestamps between start and end dates."""
        start = datetime.fromisoformat(self.config.start_date)
        end = datetime.fromisoformat(self.config.end_date)

        timestamps = []
        current = start
        while current <= end:
            timestamps.append(current)
            current += timedelta(hours=1)

        return timestamps

    def _apply_variance_to_elements(self, df: pl.DataFrame) -> pl.DataFrame:
        """Apply variance to each element's data."""
        result_dfs = []

        for element_id in self.elements:
            element_df = df.filter(pl.col("element_id") == element_id)
            element_df = self.variance_service.apply_variance(
                element_df, element_id, self.config.variance_pct
            )
            result_dfs.append(element_df)

        return pl.concat(result_dfs).sort("timestamp", "element_id")

    def _apply_anomalies_to_elements(self, df: pl.DataFrame) -> pl.DataFrame:
        """Apply configured anomalies with cost isolation.

        Phase 1: Calculate target shares for cost isolation
        Phase 2: Redistribute to target shares during anomaly windows
        Phase 3: Apply calculated telemetry multiplier to anomaly element
        """
        # Get anomaly element
        anomaly_element = self.config.get_anomaly_element()
        if not anomaly_element:
            return df  # No anomaly configured

        # Get element config
        element_cfg = self.config.parse_element_config(
            self.config.elements[anomaly_element]
        )

        # Only process spike anomalies (steps deferred to future)
        if element_cfg.anomaly_type not in ["daily_spike", "hourly_spike"]:
            return df

        # Validate cost_multiplier is provided
        if not element_cfg.cost_multiplier:
            raise ValueError(
                f"cost_multiplier required for {element_cfg.anomaly_type} "
                f"on element '{anomaly_element}'"
            )

        # Identify anomaly windows
        anomaly_timestamps = self._identify_anomaly_windows(anomaly_element)

        # Redistribute with cost isolation during anomaly windows
        # Cap telemetry spike to realistic business activity increases
        df = self._redistribute_with_cost_isolation(
            df,
            anomaly_timestamps,
            element_cfg.cost_multiplier,
            MAX_TELEMETRY_MULTIPLIER,
            anomaly_element,
        )

        return df

    def _calculate_baseline_shares(self, df: pl.DataFrame) -> Dict[str, float]:
        """Calculate each element's share of total telemetry from Pass 3 output.

        Uses average share across all timestamps to smooth out variance.

        Args:
            df: DataFrame after Pass 3 (with variance applied)

        Returns:
            Dict mapping element_id to share (e.g., {"api-search": 0.6, ...})
        """
        # Calculate total unit_value per element and shares
        element_totals = (
            df.group_by("element_id")
            .agg(pl.col("unit_value").sum().alias("total"))
        )

        # Calculate grand total
        grand_total = element_totals["total"].sum()

        # Calculate shares using vectorized operations
        element_shares = element_totals.with_columns([
            (pl.col("total") / grand_total).alias("share")
        ])

        # Convert to dict
        return dict(zip(
            element_shares["element_id"].to_list(),
            element_shares["share"].to_list()
        ))

    def _redistribute_with_cost_isolation(
        self,
        df: pl.DataFrame,
        anomaly_timestamps: Set[datetime],
        cost_multiplier: float,
        max_telemetry_multiplier: float,
        anomaly_element: str,
    ) -> pl.DataFrame:
        """Redistribute with cost isolation during anomaly windows.

        Caps anomaly element at max_telemetry_multiplier, scales down non-anomaly
        elements to maintain their baseline costs.

        Trade-off: Non-anomaly telemetry decreases during spike (counterintuitive
        but necessary for cost isolation with capped anomaly element).

        Args:
            df: DataFrame to redistribute
            anomaly_timestamps: Set of timestamps in anomaly windows
            cost_multiplier: Cost increase factor (e.g., 6.0 for 6x)
            max_telemetry_multiplier: Maximum telemetry increase (e.g., 10.0)
            anomaly_element: Element with anomaly

        Returns:
            DataFrame with capped spike and scaled-down non-anomaly elements
        """
        if not anomaly_timestamps:
            return df

        def redistribute_if_anomaly(group: pl.DataFrame) -> pl.DataFrame:
            timestamp = group["timestamp"][0]
            if timestamp not in anomaly_timestamps:
                return group  # Not in anomaly window

            # Get current values and calculate baseline shares
            total = group["unit_value"].sum()
            current_values = {}
            baseline_shares = {}
            for row_idx in range(len(group)):
                elem_id = group["element_id"][row_idx]
                value = group["unit_value"][row_idx]
                current_values[elem_id] = value
                baseline_shares[elem_id] = value / total if total > 0 else 0

            # Cap anomaly element at max_telemetry_multiplier
            anomaly_baseline = current_values[anomaly_element]
            anomaly_capped = anomaly_baseline * max_telemetry_multiplier

            # Calculate target shares for cost isolation
            target_shares = self.cost_isolation_service.calculate_target_shares(
                baseline_shares, cost_multiplier, anomaly_element
            )

            # Calculate required total telemetry given capped anomaly value
            # anomaly_capped = target_total × target_shares[anomaly_element]
            # Therefore: target_total = anomaly_capped / target_shares[anomaly_element]
            target_total = anomaly_capped / target_shares[anomaly_element]

            # Calculate target values for each element
            target_values = {}
            for element_id in current_values.keys():
                if element_id == anomaly_element:
                    target_values[element_id] = anomaly_capped
                else:
                    # Scale down non-anomaly elements to maintain cost isolation
                    target_values[element_id] = target_total * target_shares[element_id]

            # Build when/then chain for redistribution
            expr = None
            for element_id, target_value in target_values.items():
                condition = pl.col("element_id") == element_id
                value = pl.lit(target_value)

                if expr is None:
                    expr = pl.when(condition).then(value)
                else:
                    expr = expr.when(condition).then(value)

            return group.with_columns([expr.otherwise(None).alias("unit_value")])

        return df.group_by("timestamp").map_groups(redistribute_if_anomaly)

    def _apply_multiplier_during_windows(
        self,
        df: pl.DataFrame,
        element_id: str,
        anomaly_timestamps: Set[datetime],
        multiplier: float,
    ) -> pl.DataFrame:
        """Apply multiplier to specific element during anomaly windows.

        Args:
            df: DataFrame to modify
            element_id: Element to apply multiplier to
            anomaly_timestamps: Set of timestamps in anomaly windows
            multiplier: Telemetry multiplier to apply

        Returns:
            DataFrame with multiplier applied
        """
        # Create condition: element matches AND timestamp in anomaly window
        in_window = pl.col("timestamp").is_in(list(anomaly_timestamps))
        is_anomaly_element = pl.col("element_id") == element_id

        return df.with_columns([
            pl.when(in_window & is_anomaly_element)
            .then(pl.col("unit_value") * multiplier)
            .otherwise(pl.col("unit_value"))
            .alias("unit_value")
        ])

    def _identify_anomaly_windows(self, anomaly_element_id: str) -> Set[datetime]:
        """Identify timestamps that fall within anomaly windows.

        Returns set of datetime objects for efficient lookup.
        """
        # Get element config
        element_cfg = self.config.parse_element_config(
            self.config.elements.get(anomaly_element_id)
        )

        if not element_cfg.anomaly_type:
            return set()

        # Get anomaly timing
        anomaly_type = element_cfg.anomaly_type
        day, hour, period = calculate_anomaly_timing(anomaly_type)

        # Get job timeline boundaries
        start = datetime.fromisoformat(self.config.start_date)
        end = datetime.fromisoformat(self.config.end_date)

        # Calculate anomaly occurrences
        occurrence_times = calculate_anomaly_occurrences_in_job(
            start, end, period, day, hour
        )

        # Expand occurrences into full window timestamps
        anomaly_timestamps = set()
        for occurrence_time in occurrence_times:
            if "spike" in anomaly_type:
                # Spikes affect specific time windows
                if "daily" in anomaly_type:
                    # daily_spike: 3 consecutive days
                    for day_offset in range(3):
                        spike_date = occurrence_time.date() + timedelta(days=day_offset)
                        # Add all 24 hours of each spike day
                        for hour_offset in range(24):
                            timestamp = datetime.combine(spike_date, datetime.min.time()) + timedelta(hours=hour_offset)
                            anomaly_timestamps.add(timestamp)
                else:
                    # hourly_spike: 4 consecutive hours
                    for hour_offset in range(4):
                        timestamp = occurrence_time + timedelta(hours=hour_offset)
                        anomaly_timestamps.add(timestamp)
            else:
                # Step anomalies affect all timestamps after occurrence
                # But redistribution only during the step transition
                # For simplicity, don't redistribute for steps (multiplier handles it)
                pass

        return anomaly_timestamps
