"""
Reference data filtering for generation jobs.

Extracted from the existing pipeline with enhanced hierarchical filtering support.
"""

from pathlib import Path
from typing import Any, Dict, List, Tuple
import polars as pl


class ReferenceDataValidationError(Exception):
    """Raised when filtered reference data does not match filter config."""
    pass


class FilterValidationError(Exception):
    """Raised when filter structure is invalid."""
    pass


FILTER_COLUMNS = [
    "cloud_provider",
    "service", 
    "usage_family",
    "resource_type",
    "operation",
    "usage_type",
]

EXPECTED_FILTER_FIELDS = {
    "cloud_provider",
    "service",
    "usage_family", 
    "resource_type",
    "operation",
    "usage_type",
}


def apply_filters(
    ref_df: pl.DataFrame, filters: List[dict]
) -> List[Tuple[pl.DataFrame, dict]]:
    """
    Apply hierarchical filters to the reference DataFrame with indexed operations.
    Returns a list of (filtered DataFrame, filter_config) tuples.
    
    Uses optimized filtering with indexed operations and vectorized conditions
    for better performance on large reference datasets.
    """
    if ref_df.is_empty():
        return []
    
    results = []

    for filter_cfg in filters:
        cloud_provider = filter_cfg["cloud_provider"]
        service_cfgs = filter_cfg.get("services", [])

        # Optimize by pre-filtering at provider level using indexing
        provider_df = ref_df.filter(pl.col("cloud_provider") == cloud_provider)
        if provider_df.is_empty():
            continue

        # Collect all service names for vectorized filtering
        service_names = [service_cfg["name"] for service_cfg in service_cfgs]
        if not service_names:
            continue
            
        # Use vectorized IN operation for services
        services_df = provider_df.filter(pl.col("service").is_in(service_names))
        if services_df.is_empty():
            continue

        for service_cfg in service_cfgs:
            service_name = service_cfg["name"]
            service_df = services_df.filter(pl.col("service") == service_name)

            if service_df.is_empty():
                continue

            usage_family_cfgs = service_cfg.get("usage_families", [])

            # Treat both missing and explicitly empty arrays as wildcards
            if not usage_family_cfgs:
                usage_family_cfgs = [{"name": None, "resource_types": []}]  # wildcard

            # Collect usage families for vectorized filtering
            usage_family_names = [
                uf_cfg.get("name") for uf_cfg in usage_family_cfgs 
                if uf_cfg.get("name") is not None
            ]
            
            # Pre-filter usage families if any specific ones are requested
            if usage_family_names:
                usage_families_df = service_df.filter(
                    pl.col("usage_family").is_in(usage_family_names)
                )
            else:
                usage_families_df = service_df  # Wildcard case

            for usage_family_cfg in usage_family_cfgs:
                usage_family_name = usage_family_cfg.get("name")
                uf_df = (
                    usage_families_df
                    if usage_family_name is None
                    else usage_families_df.filter(pl.col("usage_family") == usage_family_name)
                )

                if uf_df.is_empty():
                    continue

                resource_type_cfgs = usage_family_cfg.get("resource_types", [])

                # Treat both missing and explicitly empty arrays as wildcards
                if not resource_type_cfgs:
                    resource_type_cfgs = [{"name": None, "operations": []}]  # wildcard

                # Collect resource types for vectorized filtering
                resource_type_names = [
                    rt_cfg.get("name") for rt_cfg in resource_type_cfgs 
                    if rt_cfg.get("name") is not None
                ]
                
                # Pre-filter resource types if any specific ones are requested
                if resource_type_names:
                    resource_types_df = uf_df.filter(
                        pl.col("resource_type").is_in(resource_type_names)
                    )
                else:
                    resource_types_df = uf_df  # Wildcard case

                for resource_type_cfg in resource_type_cfgs:
                    resource_type_name = resource_type_cfg.get("name")
                    rt_df = (
                        resource_types_df
                        if resource_type_name is None
                        else resource_types_df.filter(pl.col("resource_type") == resource_type_name)
                    )

                    if rt_df.is_empty():
                        continue

                    operation_cfgs = resource_type_cfg.get("operations", [])

                    # Treat both missing and explicitly empty arrays as wildcards
                    if not operation_cfgs:
                        operation_cfgs = [{"name": None, "usage_types": []}]  # wildcard

                    # Collect operations for vectorized filtering
                    operation_names = [
                        op_cfg.get("name") for op_cfg in operation_cfgs 
                        if op_cfg.get("name") is not None
                    ]
                    
                    # Pre-filter operations if any specific ones are requested
                    if operation_names:
                        operations_df = rt_df.filter(
                            pl.col("operation").is_in(operation_names)
                        )
                    else:
                        operations_df = rt_df  # Wildcard case

                    for operation_cfg in operation_cfgs:
                        operation_name = operation_cfg.get("name")
                        op_df = (
                            operations_df
                            if operation_name is None
                            else operations_df.filter(pl.col("operation") == operation_name)
                        )

                        if op_df.is_empty():
                            continue

                        usage_types = operation_cfg.get("usage_types", [])

                        # Only filter by usage_types if specific ones are provided
                        # Empty usage_types array acts as wildcard (matches all)
                        if usage_types:
                            final_df = op_df.filter(
                                pl.col("usage_type").is_in(usage_types)
                            )
                        else:
                            final_df = op_df

                        if not final_df.is_empty():
                            results.append(
                                (final_df, filter_cfg)
                            )  # each filtered block tied to original config

    return results


def validate_filter_structure(filters: List[dict]) -> None:
    """
    Validate that filters match the nested structure with variance profile support.
    """
    if not isinstance(filters, list):
        raise FilterValidationError("Filters must be a list of dictionaries.")

    for i, filter_cfg in enumerate(filters):
        if "cloud_provider" not in filter_cfg:
            raise FilterValidationError(
                f"Missing 'cloud_provider' field in filter #{i}"
            )

        if "services" not in filter_cfg or not isinstance(filter_cfg["services"], list):
            raise FilterValidationError(
                f"Missing or invalid 'services' list in filter #{i}"
            )

        for j, service_cfg in enumerate(filter_cfg["services"]):
            if "name" not in service_cfg:
                raise FilterValidationError(
                    f"Missing 'name' field in service #{j} of filter #{i}"
                )

            expected_service_keys = {"name", "usage_families"}
            unexpected = set(service_cfg.keys()) - expected_service_keys
            if unexpected:
                raise FilterValidationError(
                    f"Unexpected fields in service '{service_cfg['name']}': {unexpected}"
                )

            if "usage_families" not in service_cfg or not isinstance(
                service_cfg["usage_families"], list
            ):
                raise FilterValidationError(
                    f"Missing or invalid 'usage_families' list in service '{service_cfg['name']}' of filter #{i}"
                )

            for k, usage_family_cfg in enumerate(service_cfg["usage_families"]):
                if "name" not in usage_family_cfg:
                    raise FilterValidationError(
                        f"Missing 'name' field in usage_family #{k} of service '{service_cfg['name']}'"
                    )

                expected_usage_family_keys = {"name", "resource_types"}
                unexpected = set(usage_family_cfg.keys()) - expected_usage_family_keys
                if unexpected:
                    raise FilterValidationError(
                        f"Unexpected fields in service '{usage_family_cfg['name']}': {unexpected}"
                    )

                if "resource_types" not in usage_family_cfg or not isinstance(
                    usage_family_cfg["resource_types"], list
                ):
                    raise FilterValidationError(
                        f"Missing or invalid 'resource_types' list in usage_family '{usage_family_cfg['name']}' of service '{service_cfg['name']}'"
                    )

                for l, resource_type_cfg in enumerate(
                    usage_family_cfg["resource_types"]
                ):
                    if "name" not in resource_type_cfg:
                        raise FilterValidationError(
                            f"Missing 'name' field in resource_type #{l} of usage_family '{usage_family_cfg['name']}' in service '{service_cfg['name']}'"
                        )

                    expected_resource_type_keys = {"name", "operations"}
                    unexpected = (
                        set(resource_type_cfg.keys()) - expected_resource_type_keys
                    )
                    if unexpected:
                        raise FilterValidationError(
                            f"Unexpected fields in service '{resource_type_cfg['name']}': {unexpected}"
                        )

                    if "operations" not in resource_type_cfg or not isinstance(
                        resource_type_cfg["operations"], list
                    ):
                        raise FilterValidationError(
                            f"Missing or invalid 'operations' list in resource_type '{resource_type_cfg['name']}' of usage_family '{usage_family_cfg['name']}'"
                        )

                    for m, operation_cfg in enumerate(resource_type_cfg["operations"]):
                        if "name" not in operation_cfg:
                            raise FilterValidationError(
                                f"Missing 'name' field in operation #{m} of resource_type '{resource_type_cfg['name']}'"
                            )

                        expected_operation_keys = {"name", "usage_types"}
                        unexpected = set(operation_cfg.keys()) - expected_operation_keys
                        if unexpected:
                            raise FilterValidationError(
                                f"Unexpected fields in service '{operation_cfg['name']}': {unexpected}"
                            )

                        if "usage_types" not in operation_cfg or not isinstance(
                            operation_cfg["usage_types"], list
                        ):
                            raise FilterValidationError(
                                f"Missing or invalid 'usage_types' list in operation '{operation_cfg['name']}' of resource_type '{resource_type_cfg['name']}'"
                            )

                        for usage_type in operation_cfg["usage_types"]:
                            if not isinstance(usage_type, str):
                                raise FilterValidationError(
                                    f"Each usage_type in operation '{operation_cfg['name']}' must be a string"
                                )

        # Validate optional variance profile fields at filter level
        if "variance_profile" in filter_cfg:
            variance_profile = filter_cfg["variance_profile"]
            if not isinstance(variance_profile, str):
                raise FilterValidationError(
                    f"'variance_profile' in filter #{i} must be a string"
                )

        if "anomaly_profile" in filter_cfg:
            anomaly_profile = filter_cfg["anomaly_profile"]
            # Support both string and dict formats for anomaly_profile
            if isinstance(anomaly_profile, str):
                # String format: "daily_spike"
                pass  # Valid
            elif isinstance(anomaly_profile, dict):
                # Dict format: {"type": "daily_spike", "intensity": 1.5, ...}
                if "type" not in anomaly_profile:
                    raise FilterValidationError(
                        f"'anomaly_profile' dict in filter #{i} must contain 'type' field"
                    )
            elif anomaly_profile != "" and anomaly_profile is not None:
                # Only raise error if it's not empty string or None
                raise FilterValidationError(
                    f"'anomaly_profile' in filter #{i} must be a string or dict, got {type(anomaly_profile).__name__}"
                )


def validate_filtered_groups(ref_df: pl.DataFrame, filters: List[dict]) -> dict:
    """
    Validate that the reference data contains all expected services, usage_families, 
    resource_types, operations, and usage_types, scoped correctly according to the 
    nested structure.
    """
    diffs = {}

    for filter_cfg in filters:
        cloud_provider = filter_cfg["cloud_provider"]
        provider_df = ref_df.filter(pl.col("cloud_provider") == cloud_provider)

        for service_cfg in filter_cfg.get("services", []):
            service_name = service_cfg["name"]
            service_df = provider_df.filter(pl.col("service") == service_name)

            service_diff = {}

            if service_df.is_empty():
                service_diff["missing_service"] = True
                diffs.setdefault(cloud_provider, {})[service_name] = service_diff
                continue

            for usage_family_cfg in service_cfg.get("usage_families", []):
                uf_name = usage_family_cfg["name"]
                uf_df = service_df.filter(pl.col("usage_family") == uf_name)

                uf_diff = {}

                if uf_df.is_empty():
                    uf_diff["missing_usage_family"] = True
                    service_diff.setdefault(uf_name, uf_diff)
                    continue

                for resource_type_cfg in usage_family_cfg.get("resource_types", []):
                    rt_name = resource_type_cfg["name"]
                    rt_df = uf_df.filter(pl.col("resource_type") == rt_name)

                    rt_diff = {}

                    if rt_df.is_empty():
                        rt_diff["missing_resource_type"] = True
                        service_diff.setdefault(uf_name, {}).setdefault(
                            rt_name, rt_diff
                        )
                        continue

                    for operation_cfg in resource_type_cfg.get("operations", []):
                        op_name = operation_cfg["name"]
                        op_df = rt_df.filter(pl.col("operation") == op_name)

                        op_diff = {}

                        if op_df.is_empty():
                            op_diff["missing_operation"] = True
                            service_diff.setdefault(uf_name, {}).setdefault(
                                rt_name, {}
                            ).setdefault(op_name, op_diff)
                            continue

                        expected_usage_types = set(operation_cfg.get("usage_types", []))
                        if expected_usage_types:
                            actual_usage_types = set(
                                op_df["usage_type"].drop_nulls().to_list()
                            )
                            missing_usage_types = (
                                expected_usage_types - actual_usage_types
                            )
                            if missing_usage_types:
                                op_diff["missing_usage_types"] = list(
                                    missing_usage_types
                                )
                                service_diff.setdefault(uf_name, {}).setdefault(
                                    rt_name, {}
                                ).setdefault(op_name, op_diff)

            if service_diff:
                diffs.setdefault(cloud_provider, {})[service_name] = service_diff

    return diffs


def load_reference_data_for_job(
    ref_data_path: Path,
    cloud_provider: str,
    services: List[str]
) -> pl.DataFrame:
    """
    Load reference data optimized for a specific job using lazy evaluation.
    
    Filters reference data by provider and services using lazy operations
    to improve memory efficiency and processing speed.
    """
    # Load reference data with lazy evaluation for better memory efficiency
    filtered_ref_df = (
        pl.scan_parquet(ref_data_path)
        .filter(
            (pl.col("cloud_provider") == cloud_provider) &
            (pl.col("service").is_in(services))
        )
        .collect()  # Only collect after applying filters
    )
    
    return filtered_ref_df


def estimate_record_count_after_filtering(
    ref_df: pl.DataFrame,
    filters: List[dict],
    resource_count: int,
    hours_in_range: int
) -> int:
    """
    Estimate the total number of records that will be generated after filtering.
    
    This is called after reference data filtering to provide accurate estimates
    for memory management and job splitting decisions.
    """
    # Apply filters to get actual matching reference data
    filtered_results = apply_filters(ref_df, filters)
    
    # Count total reference matches
    total_reference_matches = sum(len(filtered_df) for filtered_df, _ in filtered_results)
    
    # Calculate total records: reference matches × resource_count × hours
    total_records = total_reference_matches * resource_count * hours_in_range
    
    return total_records