"""Defines the Metric class for evaluating data quality metrics."""

from enum import StrEnum

import numpy as np
from ctapipe.core import Component, TelescopeComponent, non_abstract_children
from ctapipe.core.traits import Bool, List, Unicode
from ctapipe.core.traits import Enum as TraitEnum
from ctapipe.monitoring.aggregator import BaseAggregator
from traitlets.config import Config

from .criterion import BaseQualityCriterion, BaseTelescopeQualityCriterion
from .descriptor import TrivialDescriptor

__all__ = [
    "ArrayElement",
    "DataCategory",
    "Metric",
    "get_available_aggregators",
    "get_available_criteria",
]


def get_available_aggregators():
    """
    Get all available aggregator/descriptor classes (non-abstract subclasses of BaseAggregator).

    Returns
    -------
    dict
        Dictionary mapping class names to class objects for all available aggregators
        from both ctapipe and qualpipe.
    """
    # Use ctapipe's utility to find all non-abstract subclasses
    aggregators = non_abstract_children(BaseAggregator)
    return {cls.__name__: cls for cls in aggregators}


def get_available_criteria():
    """
    Get all available quality criteria classes (non-abstract subclasses of Component or TelescopeComponent).

    Returns
    -------
    dict
        Dictionary mapping class names to class objects for all available criteria
        from qualpipe. Includes both regular and telescope-specific criteria.
    """
    # Use ctapipe's utility to find all non-abstract subclasses
    criteria = non_abstract_children(BaseQualityCriterion)

    # Also get telescope-specific criteria
    tel_criteria = non_abstract_children(BaseTelescopeQualityCriterion)

    # Combine both into a single dict (telescope criteria may overlap, but that's fine)
    result = {cls.__name__: cls for cls in criteria}
    result.update({cls.__name__: cls for cls in tel_criteria})

    return result


class ArrayElement(StrEnum):
    """Enumeration of array elements."""

    DUMMY = "DUMMY"
    SUBARRAY = "SUBARRAY"
    LST = "LST"
    MST = "MST"
    SST = "SST"
    FRAM = "FRAM"
    LIDAR = "LIDAR"
    WEATHER_STATION = "WEATHER_STATION"


class DataCategory(StrEnum):
    """Enumeration of data categories."""

    DL0_EVENT = "DL0_EVENT"
    DL1_EVENT = "DL1_EVENT"
    DL2_EVENT = "DL2_EVENT"
    DL3_EVENT = "DL3_EVENT"
    DL0_SERVICE = "DL0_SERVICE"
    DL1_SERVICE = "DL1_SERVICE"
    DL2_SERVICE = "DL2_SERVICE"
    DL3_SERVICE = "DL3_SERVICE"
    DL0_MONITORING = "DL0_MONITORING"
    DL1_MONITORING = "DL1_MONITORING"
    DL2_MONITORING = "DL2_MONITORING"
    DL3_MONITORING = "DL3_MONITORING"


class Metric(Component):
    """Defines a data quality metric that applies a descriptor/aggregator to input data and evaluates quality criteria."""

    name = Unicode(help="Metric name").tag(config=True)

    input_source = Unicode(help="HDF5 table path").tag(config=True)

    input_source_subtype = Unicode(
        None,
        allow_none=True,
        help="Optional subtype of input source (e.g., 'flatfield_image') for more specific data selection.",
    ).tag(config=True)

    table_column = Unicode(
        None,
        allow_none=True,
        help="Which column from the input table to use for the metric calculation. ",
    ).tag(config=True)

    array_element = TraitEnum(
        ArrayElement,
        help="Array element",
    ).tag(config=True)
    telescope_specific = Bool(
        False, help="Flag to indicate if this is a telescope-specific metric"
    ).tag(config=True)
    descriptor = Unicode(
        None,
        allow_none=True,
        help="Name of the aggregator/descriptor to use (e.g., PlainAggregator, PercentageDescriptor)",
    ).tag(config=True)
    descriptor_column = Unicode(
        None,
        allow_none=True,
        help="Which column from descriptor output to use for criteria evaluation. "
        "For ctapipe aggregators: 'mean', 'median', or 'std'. "
        "For custom descriptors: varies (e.g., 'percentage', 'data'). "
        "If None, uses first non-metadata column.",
    ).tag(config=True)
    criteria = List(help="List of configurations for quality criteria").tag(config=True)

    def __init__(self, **kwargs):
        self.subarray = kwargs.pop("subarray", None)
        super().__init__(**kwargs)

        # Load descriptor/aggregator
        # Note: Aggregators are NOT telescope-specific, they just process tables
        if self.descriptor is None or self.descriptor == "TrivialDescriptor":
            # No descriptor specified or TrivialDescriptor explicitly requested
            # This returns the input table as-is (entire dataset as one chunk)
            self._descriptor = TrivialDescriptor()
        else:
            # Use ctapipe Component registry for other descriptors/aggregators
            self._descriptor = Component.from_name(self.descriptor, parent=self)

        # Load quality criteria
        # Criteria can be telescope-specific if needed
        if self.telescope_specific:
            base_component_class = TelescopeComponent
            kwargs_dict = {"subarray": self.subarray}
        else:
            base_component_class = Component
            kwargs_dict = {}

        self._criteria = [
            base_component_class.from_name(
                next(iter(criterion)), config=Config(criterion), **kwargs_dict
            )
            for criterion in self.criteria
        ]

        # Validate that there are no duplicate criterion types
        # (since column names are based on class name, duplicates would overwrite)
        criterion_types = [type(c).__name__ for c in self._criteria]
        duplicates = [t for t in set(criterion_types) if criterion_types.count(t) > 1]
        if duplicates:
            raise ValueError(
                f"Metric '{self.name}' has multiple criteria of the same type: {duplicates}. "
                "Each criterion type can only be used once per metric, as they share "
                "the same column name in the output table. Consider using different "
                "criterion types or splitting into separate metrics."
            )

        self.log.debug("Metric %s initialized", self.name)
        self.log.debug("Metric's descriptor: %s", self.descriptor)
        self.log.debug("Metric's criteria:")
        for criterion in self._criteria:
            self.log.debug(criterion)

    def __call__(self, table, tel_id=None):
        """
        Compute the metric value using the descriptor/aggregator and evaluate quality criteria.

        Parameters
        ----------
        table : astropy.table.Table
            Input table with data. Must have 'time_mono', 'event_id', and data columns.
        tel_id : int, optional
            Telescope ID for telescope-specific metrics (used only for criteria evaluation)

        Returns
        -------
        astropy.table.Table
            Results table with one row per time window, including:
            - Time window information (time_start, time_end, event_id_start, event_id_end, n_events)
            - Descriptor/aggregator statistics (e.g., mean, median, std)
            - Criterion pass/fail columns ({CriterionName}_pass)
            - Overall pass/fail (overall_pass) and pass rate (pass_rate)
            - Metadata with metric configuration
        """
        if self.telescope_specific and tel_id is None:
            raise ValueError("Telescope-specific metric requires a telescope ID")

        # Apply descriptor/aggregator
        # Aggregator always returns a table with time windows and computed statistics
        if isinstance(self._descriptor, TrivialDescriptor):
            col_name = self.descriptor_column if self.descriptor_column else "data"
        else:
            col_name = self._get_data_column_name(table)

        result_table = self._descriptor(table, col_name=col_name)

        # Determine which column to use for criteria evaluation
        statistic_column = self._get_statistic_column(result_table)

        # Apply criteria to each time window (modifies result_table in place)
        self._apply_criteria_to_time_series(result_table, statistic_column, tel_id)

        # Get the current configuration
        config_to_add = self.get_current_config()

        # Add safe metric configuration to table metadata (HDF5-compatible types only)
        # FIXME: Remove this workaround when ctapipe supports complex metadata types
        def make_hdf5_safe(obj, max_depth=125):
            """Recursively convert objects to HDF5-safe types."""
            if max_depth <= 0:
                return str(obj)  # Convert to string if too deep

            if obj is None or isinstance(obj, bool | int | float | str):
                return obj
            elif isinstance(obj, list | tuple):
                return [make_hdf5_safe(item, max_depth - 1) for item in obj]
            elif isinstance(obj, dict):
                return {
                    key: make_hdf5_safe(value, max_depth - 1)
                    for key, value in obj.items()
                }
            else:
                # Convert complex objects to strings
                return str(obj)

        safe_config = make_hdf5_safe(config_to_add)
        result_table.meta["metric_config"] = safe_config
        result_table.meta["n_time_windows"] = len(result_table)

        return result_table

    def _get_data_column_name(self, table):
        """
        Determine the data column name from input table.

        By convention, MetricEvaluator should rename the target column to 'data'.
        """
        if "data" in table.colnames:
            return "data"
        else:
            # Fallback: use the column name from
            return self.table_column

    def _get_statistic_column(self, result_table):
        """
        Determine which column contains the statistic for criteria evaluation.

        Returns
        -------
        str
            Name of the column to use for criteria
        """
        if self.descriptor_column is not None:
            # User explicitly specified which column
            if self.descriptor_column not in result_table.colnames:
                raise ValueError(
                    f"Specified descriptor_column '{self.descriptor_column}' "
                    f"not found in result table. Available: {result_table.colnames}"
                )
            return self.descriptor_column

        # Auto-detect: use first non-metadata column
        metadata_cols = {
            "time_start",
            "time_end",
            "event_id_start",
            "event_id_end",
            "n_events",
        }
        statistic_cols = [
            col for col in result_table.colnames if col not in metadata_cols
        ]

        if not statistic_cols:
            raise ValueError("No statistic columns found in descriptor output")

        # Use first statistic column
        return statistic_cols[0]

    def _apply_criteria_to_time_series(self, result_table, statistic_column, tel_id):
        """
        Apply all criteria to the time series of statistic values and add columns to result table.

        Parameters
        ----------
        result_table : astropy.table.Table
            Table with time windows and statistics (modified in place)
        statistic_column : str
            Name of column containing values to evaluate
        tel_id : int, optional
            Telescope ID for telescope-specific criteria

        Returns
        -------
        None
            Modifies result_table in place by adding criterion pass/fail columns
        """
        for criterion in self._criteria:
            criterion_name = criterion.__class__.__name__
            criterion_results = []

            # Apply criterion to each time window
            for row in result_table:
                value = row[statistic_column]
                result = criterion(value, tel_id if self.telescope_specific else None)
                criterion_results.append(result)

            # Add pass/fail column for this criterion
            result_table[f"{criterion_name}_pass"] = criterion_results

        # Add overall_pass column (True only if ALL criteria passed)
        # Create boolean mask for each criterion
        criterion_columns = [
            result_table[f"{criterion.__class__.__name__}_pass"]
            for criterion in self._criteria
        ]
        # All criteria must pass for overall_pass to be True
        result_table["overall_pass"] = np.all(criterion_columns, axis=0)

        # Add pass_rate column (fraction of criteria that passed per time window)
        result_table["pass_rate"] = np.mean(criterion_columns, axis=0)
