"""
ConfigPartitioner Service - Break pipeline configs into minimal processing units.

Primary Responsibility: Convert complex pipeline configurations into sequential,
memory-bounded processing parts for immediate write pattern.

Key Features:
- Sequential config part generation
- Provider + service + bounded reference data per part
- Deterministic partitioning for consistent processing
- Memory-bounded part sizing
- Predictable processing order
"""

import hashlib
from collections.abc import Iterator
from dataclasses import dataclass
from typing import (
    Any,
    Dict,
    List,
    Optional,
)

import polars as pl


@dataclass
class ConfigPart:
    """
    A minimal processing unit for generate stage.

    Represents a single provider + service combination with all
    necessary configuration and bounded reference data subset.
    """

    # Core identifiers
    provider: str
    service: str

    # Configuration data
    reference_data_path: str
    start_date: str
    end_date: str
    total_cost: float
    resource_count: int

    # Service-specific configuration
    service_config: Dict[str, Any]

    # Variance and anomaly profiles
    variance_profile: Optional[str] = None
    anomaly_profile: Optional[str] = None

    # Processing metadata
    part_id: str = ""
    estimated_memory_mb: float = 0.0

    def __post_init__(self):
        """Generate deterministic part ID for consistent processing."""
        if not self.part_id:
            # Create deterministic hash from core identifiers
            hash_input = f"{self.provider}:{self.service}:{self.reference_data_path}"
            self.part_id = hashlib.md5(hash_input.encode()).hexdigest()[:12]


class ConfigPartitioner:
    """
    Service for breaking pipeline configurations into minimal processing units.

    Transforms complex pipeline configs with multiple providers and services
    into sequential, memory-bounded config parts suitable for immediate write pattern.

    Features:
    - Deterministic partitioning for consistent processing
    - Memory estimation for each config part
    - Provider/service isolation for clean processing boundaries
    - Reference data subset coordination
    """

    def __init__(self, max_part_memory_mb: float = 500.0):
        """
        Initialize ConfigPartitioner.

        Args:
            max_part_memory_mb: Maximum memory per config part in MB
        """
        self.max_part_memory_mb = max_part_memory_mb
        self._reference_data_cache: Dict[str, pl.DataFrame] = {}

    def get_config_parts(self, pipeline_config: dict) -> Iterator[ConfigPart]:
        """
        Generate sequential config parts from pipeline configuration.

        Breaks down complex pipeline config into minimal processing units:
        1. Extract each job from the generate stage
        2. For each job, create parts for each provider/service combination
        3. Estimate memory usage and validate bounds
        4. Yield parts in deterministic order for consistent processing

        Args:
            pipeline_config: Complete pipeline configuration from YAML

        Yields:
            ConfigPart: Minimal processing unit ready for immediate write
        """
        generate_config = pipeline_config.get("stages", {}).get("generate", [])

        if not generate_config:
            return

        # Process each job in the generate stage
        for job_index, job_config in enumerate(generate_config):
            yield from self._create_parts_for_job(job_config, job_index)

    def _create_parts_for_job(
        self, job_config: dict, job_index: int
    ) -> Iterator[ConfigPart]:
        """
        Create config parts for a single job configuration.

        Args:
            job_config: Single job from generate stage configuration
            job_index: Index of job for deterministic ordering

        Yields:
            ConfigPart: Parts for each provider/service combination
        """
        # Extract job-level configuration
        reference_data_path = job_config.get("reference_data", "")
        start_date = job_config.get("start_date", "")
        end_date = job_config.get("end_date", "")
        total_cost = float(job_config.get("total_cost", 0.0))
        resource_count = int(job_config.get("resource_count", 1))

        filters = job_config.get("filters", [])

        # Process each filter (provider configuration)
        for filter_index, filter_config in enumerate(filters):
            provider = filter_config.get("cloud_provider", "")
            variance_profile = filter_config.get("variance_profile")
            anomaly_profile = filter_config.get("anomaly_profile")

            services = filter_config.get("services", [])

            # Calculate cost per service (divide total cost among services in this filter)
            num_services = len(services) if services else 1
            cost_per_service = total_cost / num_services

            # Create config part for each service
            for service_index, service_config in enumerate(services):
                service_name = service_config.get("name", "")

                # Create config part for this provider/service combination
                config_part = ConfigPart(
                    provider=provider,
                    service=service_name,
                    reference_data_path=reference_data_path,
                    start_date=start_date,
                    end_date=end_date,
                    total_cost=cost_per_service,  # Use divided cost, not full cost
                    resource_count=resource_count,
                    service_config=service_config,
                    variance_profile=variance_profile,
                    anomaly_profile=anomaly_profile,
                )

                # Estimate memory usage for this part
                config_part.estimated_memory_mb = self.estimate_part_memory(config_part)

                # Validate memory bounds
                if config_part.estimated_memory_mb > self.max_part_memory_mb:
                    # For now, log warning but continue processing
                    # In future, could implement part splitting
                    pass

                yield config_part

    def estimate_part_memory(self, config_part: ConfigPart) -> float:
        """
        Estimate memory usage for processing a config part.

        Provides rough memory estimate based on:
        - Resource count and time range
        - Reference data subset size
        - Historical processing patterns

        Args:
            config_part: Config part to estimate

        Returns:
            Estimated memory usage in MB
        """
        # Base memory for DataFrame operations
        base_memory_mb = 50.0

        # Memory per resource (roughly 1KB per resource per month)
        resource_memory_mb = (
            config_part.resource_count * self._estimate_month_count(config_part)
        ) / 1024

        # Reference data subset memory (estimated)
        ref_data_memory_mb = self._estimate_reference_data_memory(config_part)

        # Time series generation memory (depends on date range)
        time_series_memory_mb = self._estimate_time_series_memory(config_part)

        total_memory_mb = (
            base_memory_mb
            + resource_memory_mb
            + ref_data_memory_mb
            + time_series_memory_mb
        )

        return total_memory_mb

    def _estimate_month_count(self, config_part: ConfigPart) -> int:
        """Estimate number of months in date range."""
        try:
            from datetime import datetime

            start = datetime.strptime(config_part.start_date, "%Y-%m-%d")
            end = datetime.strptime(config_part.end_date, "%Y-%m-%d")
            months = (end.year - start.year) * 12 + (end.month - start.month) + 1
            return max(1, months)
        except Exception:
            return 12  # Default to 1 year

    def _estimate_reference_data_memory(self, config_part: ConfigPart) -> float:
        """
        Estimate memory for reference data subset.

        Args:
            config_part: Config part to estimate

        Returns:
            Estimated memory in MB for reference data subset
        """
        # If we have cached reference data, use actual size
        if config_part.reference_data_path in self._reference_data_cache:
            ref_df = self._reference_data_cache[config_part.reference_data_path]

            # Filter for this provider/service to get subset size
            filtered_df = ref_df.filter(
                (pl.col("cloud_provider") == config_part.provider)
                & (pl.col("service") == config_part.service)
            )

            # Rough memory estimate: 1KB per row
            return (filtered_df.height * 1024) / (1024 * 1024)  # Convert to MB

        # Default estimate: 10MB for reference data subset
        return 10.0

    def _estimate_time_series_memory(self, config_part: ConfigPart) -> float:
        """
        Estimate memory for time series generation.

        Args:
            config_part: Config part to estimate

        Returns:
            Estimated memory in MB for time series data
        """
        # Time series memory depends on date range and resource count
        month_count = self._estimate_month_count(config_part)
        hours_per_month = 730  # Average hours per month
        total_hours = month_count * hours_per_month

        # Memory per hour per resource (roughly 500 bytes)
        memory_bytes = total_hours * config_part.resource_count * 500

        return memory_bytes / (1024 * 1024)  # Convert to MB

    def cache_reference_data(
        self, reference_data_path: str, ref_df: pl.DataFrame
    ) -> None:
        """
        Cache reference data for memory estimation.

        Args:
            reference_data_path: Path to reference data
            ref_df: Reference DataFrame to cache
        """
        self._reference_data_cache[reference_data_path] = ref_df

    def get_part_statistics(self, pipeline_config: dict) -> Dict[str, Any]:
        """
        Get statistics about config partitioning.

        Useful for understanding memory requirements and processing
        characteristics before execution.

        Args:
            pipeline_config: Complete pipeline configuration

        Returns:
            Dictionary with partitioning statistics
        """
        parts = list(self.get_config_parts(pipeline_config))

        if not parts:
            return {
                "total_parts": 0,
                "estimated_total_memory_mb": 0.0,
                "providers": [],
                "services": [],
            }

        providers = set(part.provider for part in parts)
        services = set(part.service for part in parts)
        total_memory = sum(part.estimated_memory_mb for part in parts)
        max_memory = max(part.estimated_memory_mb for part in parts)

        return {
            "total_parts": len(parts),
            "estimated_total_memory_mb": total_memory,
            "estimated_max_part_memory_mb": max_memory,
            "providers": sorted(providers),
            "services": sorted(services),
            "part_details": [
                {
                    "part_id": part.part_id,
                    "provider": part.provider,
                    "service": part.service,
                    "estimated_memory_mb": part.estimated_memory_mb,
                }
                for part in parts
            ],
        }

    def validate_config_parts(self, pipeline_config: dict) -> List[str]:
        """
        Validate config parts for potential issues.

        Args:
            pipeline_config: Complete pipeline configuration

        Returns:
            List of validation error messages (empty if valid)
        """
        errors = []
        parts = list(self.get_config_parts(pipeline_config))

        if not parts:
            errors.append("No config parts generated from pipeline configuration")
            return errors

        # Check for oversized parts
        oversized_parts = [
            part for part in parts if part.estimated_memory_mb > self.max_part_memory_mb
        ]

        for part in oversized_parts:
            errors.append(
                f"Config part {part.part_id} ({part.provider}/{part.service}) "
                f"estimated memory {part.estimated_memory_mb:.1f}MB exceeds "
                f"limit {self.max_part_memory_mb}MB"
            )

        # Check for missing required configuration
        for part in parts:
            if not part.provider:
                errors.append(f"Config part {part.part_id} missing provider")
            if not part.service:
                errors.append(f"Config part {part.part_id} missing service")
            if not part.reference_data_path:
                errors.append(f"Config part {part.part_id} missing reference_data_path")

        return errors


def create_config_partitioner(max_memory_mb: float = 500.0) -> ConfigPartitioner:
    """
    Factory function to create ConfigPartitioner with standard settings.

    Args:
        max_memory_mb: Maximum memory per config part in MB

    Returns:
        Configured ConfigPartitioner instance
    """
    return ConfigPartitioner(max_part_memory_mb=max_memory_mb)
