from __future__ import annotations

import logging
from collections import defaultdict
from copy import deepcopy
from typing import Dict, List, Set, Optional, Sequence, Tuple, FrozenSet

from metricflow.errors.errors import (
    DuplicateMetricError,
    MetricNotFoundError,
    NonExistentMeasureError,
    InvalidDataSourceError,
)
from metricflow.instances import DataSourceReference, DataSourceElementReference
from metricflow.aggregation_properties import AggregationType
from metricflow.model.objects.data_source import DataSource, DataSourceOrigin
from metricflow.model.objects.elements.dimension import Dimension
from metricflow.model.objects.elements.identifier import Identifier
from metricflow.model.objects.elements.measure import Measure
from metricflow.model.objects.metric import Metric, MetricType
from metricflow.model.objects.user_configured_model import UserConfiguredModel
from metricflow.model.semantics.data_source_container import PydanticDataSourceContainer
from metricflow.model.semantics.element_group import ElementGrouper
from metricflow.model.semantics.linkable_spec_resolver import (
    ValidLinkableSpecResolver,
    LinkableElementProperties,
)
from metricflow.model.spec_converters import MeasureConverter, WhereConstraintConverter
from metricflow.references import (
    DimensionReference,
    IdentifierReference,
    LinkableElementReference,
    MeasureReference,
    MetricReference,
    TimeDimensionReference,
)
from metricflow.specs import (
    LinkableInstanceSpec,
    MeasureSpec,
    MetricInputMeasureSpec,
    NonAdditiveDimensionSpec,
)

logger = logging.getLogger(__name__)

MAX_JOIN_HOPS = 2


class MetricSemantics:  # noqa: D
    def __init__(  # noqa: D
        self, user_configured_model: UserConfiguredModel, data_source_semantics: DataSourceSemantics
    ) -> None:
        self._user_configured_model = user_configured_model
        self._metrics: Dict[MetricReference, Metric] = {}
        self._data_source_semantics = data_source_semantics

        # Dict from the name of the metric to the hash.
        self._metric_hashes: Dict[MetricReference, str] = {}

        for metric in self._user_configured_model.metrics:
            self.add_metric(metric)

        self._linkable_spec_resolver = ValidLinkableSpecResolver(
            user_configured_model=self._user_configured_model,
            max_identifier_links=MAX_JOIN_HOPS,
        )

    def element_specs_for_metrics(
        self,
        metric_references: List[MetricReference],
        with_any_property: FrozenSet[LinkableElementProperties] = LinkableElementProperties.all_properties(),
        without_any_property: FrozenSet[LinkableElementProperties] = frozenset(),
    ) -> List[LinkableInstanceSpec]:
        """Dimensions common to all metrics requested (intersection)"""

        all_linkable_specs = self._linkable_spec_resolver.get_linkable_elements_for_metrics(
            metric_references=metric_references,
            with_any_of=with_any_property,
            without_any_of=without_any_property,
        ).as_spec_set

        return sorted(all_linkable_specs.as_tuple, key=lambda x: x.qualified_name)

    def get_metrics(self, metric_references: List[MetricReference]) -> List[Metric]:  # noqa: D
        res = []
        for metric_reference in metric_references:
            if metric_reference not in self._metrics:
                raise MetricNotFoundError(
                    f"Unable to find metric `{metric_reference}`. Perhaps it has not been registered"
                )
            res.append(self._metrics[metric_reference])

        return res

    @property
    def metric_references(self) -> List[MetricReference]:  # noqa: D
        return list(self._metrics.keys())

    def get_metric(self, metric_reference: MetricReference) -> Metric:  # noqa:D
        if metric_reference not in self._metrics:
            raise MetricNotFoundError(f"Unable to find metric `{metric_reference}`. Perhaps it has not been registered")
        return self._metrics[metric_reference]

    def add_metric(self, metric: Metric) -> None:
        """Add metric, validating presence of required measures"""
        metric_reference = MetricReference(element_name=metric.name)
        if metric_reference in self._metrics:
            raise DuplicateMetricError(f"Metric `{metric.name}` has already been registered")
        for measure_reference in metric.measure_references:
            if measure_reference not in self._data_source_semantics.measure_references:
                raise NonExistentMeasureError(
                    f"Metric `{metric.name}` references measure `{measure_reference}` which has not been registered"
                )
        self._metrics[metric_reference] = metric
        self._metric_hashes[metric_reference] = metric.definition_hash

    @property
    def valid_hashes(self) -> Set[str]:
        """Return all of the hashes of the metric definitions."""
        return set(self._metric_hashes.values())

    def measures_for_metric(self, metric_reference: MetricReference) -> Tuple[MetricInputMeasureSpec, ...]:
        """Return the measure specs required to compute the metric."""
        metric = self.get_metric(metric_reference)
        input_measure_specs: List[MetricInputMeasureSpec] = []

        for input_measure in metric.input_measures:
            spec_constraint = (
                WhereConstraintConverter.convert_to_spec_where_constraint(
                    data_source_semantics=self._data_source_semantics,
                    where_constraint=input_measure.constraint,
                )
                if input_measure.constraint is not None
                else None
            )
            measure_spec = MeasureSpec(
                element_name=input_measure.name,
                non_additive_dimension_spec=self._data_source_semantics.non_additive_dimension_specs_by_measure.get(
                    input_measure.measure_reference
                ),
            )
            spec = MetricInputMeasureSpec(
                measure_spec=measure_spec,
                constraint=spec_constraint,
                alias=input_measure.alias,
            )
            input_measure_specs.append(spec)

        return tuple(input_measure_specs)

    def contains_cumulative_metric(self, metric_references: Sequence[MetricReference]) -> bool:
        """Returns true if any of the specs correspond to a cumulative metric."""
        for metric_reference in metric_references:
            if self.get_metric(metric_reference).type == MetricType.CUMULATIVE:
                return True
        return False


class DataSourceSemantics:
    """Tracks semantic information for data source held in a set of DataSourceContainers

    This implements both the DataSourceSemanticsAccessors protocol, the interface type we use throughout the codebase.
    That interface prevents unwanted calls to methods for adding data sources to the container.
    """

    def __init__(  # noqa: D
        self,
        model: UserConfiguredModel,
        configured_data_source_container: PydanticDataSourceContainer,
    ) -> None:
        self._model = model
        self._measure_index: Dict[MeasureReference, List[DataSource]] = defaultdict(list)
        self._measure_aggs: Dict[
            MeasureReference, AggregationType
        ] = {}  # maps measures to their one consistent aggregation
        self._measure_agg_time_dimension: Dict[MeasureReference, TimeDimensionReference] = {}
        self._measure_non_additive_dimension_specs: Dict[MeasureReference, NonAdditiveDimensionSpec] = {}
        self._dimension_index: Dict[DimensionReference, List[DataSource]] = defaultdict(list)
        self._linkable_reference_index: Dict[LinkableElementReference, List[DataSource]] = defaultdict(list)
        self._entity_index: Dict[Optional[str], List[DataSource]] = defaultdict(list)
        self._identifier_ref_to_entity: Dict[IdentifierReference, Optional[str]] = {}
        self._data_source_names: Set[str] = set()

        self._configured_data_source_container = configured_data_source_container
        self._data_source_to_aggregation_time_dimensions: Dict[
            DataSourceReference, ElementGrouper[TimeDimensionReference, MeasureSpec]
        ] = {}

        # Add semantic tracking for data sources from configured_data_source_container
        for data_source in self._configured_data_source_container.values():
            assert isinstance(data_source, DataSource)
            self.add_configured_data_source(data_source)

    def add_configured_data_source(self, data_source: DataSource) -> None:
        """Dont use this unless you mean it (ie in tests). The configured data sources are supposed to be static"""
        self._configured_data_source_container._put(data_source)
        self._add_data_source(data_source)

    def get_dimension_references(self) -> List[DimensionReference]:  # noqa: D
        return list(self._dimension_index.keys())

    def get_dimension(
        self, dimension_reference: DimensionReference, origin: Optional[DataSourceOrigin] = None
    ) -> Dimension:
        """Retrieves a full dimension object by name"""
        for dimension_source in self._dimension_index[dimension_reference]:
            if origin and dimension_source.origin != origin:
                continue
            dimension = dimension_source.get_dimension(dimension_reference)
            # find the data source that has the requested dimension by the requested identifier

            return deepcopy(dimension)

        raise ValueError(
            f"Could not find dimension with name ({dimension_reference.element_name}) in configured data sources"
        )

    def get_time_dimension(self, time_dimension_reference: TimeDimensionReference) -> Dimension:
        """Retrieves a full dimension object by name"""
        dimension_reference = time_dimension_reference.dimension_reference()

        if dimension_reference not in self._dimension_index:
            raise ValueError(
                f"Could not find dimension with name ({dimension_reference.element_name}) in configured data sources"
            )

        for dimension_source in self._dimension_index[dimension_reference]:
            dimension = dimension_source.get_dimension(dimension_reference)
            # TODO: Unclear if the deepcopy is necessary.
            return deepcopy(dimension)

        assert False, f"{time_dimension_reference} should have been in the dimension index"

    @property
    def measure_references(self) -> List[MeasureReference]:  # noqa: D
        return list(self._measure_index.keys())

    @property
    def non_additive_dimension_specs_by_measure(self) -> Dict[MeasureReference, NonAdditiveDimensionSpec]:  # noqa: D
        return self._measure_non_additive_dimension_specs

    def get_measure(self, measure_reference: MeasureReference) -> Measure:  # noqa: D
        if measure_reference not in self._measure_index:
            raise ValueError(f"Could not find measure with name ({measure_reference}) in configured data sources")

        assert len(self._measure_index[measure_reference]) >= 1
        # Measures should be consistent across data sources, so just use the first one.
        return list(self._measure_index[measure_reference])[0].get_measure(measure_reference)

    def get_identifier_references(self) -> List[IdentifierReference]:  # noqa: D
        return list(self._identifier_ref_to_entity.keys())

    # DSC interface
    def get_data_sources_for_measure(self, measure_reference: MeasureReference) -> List[DataSource]:  # noqa: D
        return self._measure_index[measure_reference]

    def get_agg_time_dimension_for_measure(  # noqa: D
        self, measure_reference: MeasureReference
    ) -> TimeDimensionReference:
        return self._measure_agg_time_dimension[measure_reference]

    def get_identifier_in_data_source(self, ref: DataSourceElementReference) -> Optional[Identifier]:  # Noqa: d
        data_source = self.get(ref.data_source_name)
        if not data_source:
            return None

        for identifier in data_source.identifiers:
            if identifier.reference.element_name == ref.element_name:
                return identifier

        return None

    def get(self, data_source_name: str) -> Optional[DataSource]:  # noqa: D
        if data_source_name in self._configured_data_source_container:
            data_source = self._configured_data_source_container.get(data_source_name)
            assert isinstance(data_source, DataSource)
            return data_source

        return None

    def get_by_reference(self, data_source_reference: DataSourceReference) -> Optional[DataSource]:  # noqa: D
        return self.get(data_source_reference.data_source_name)

    def _add_data_source(
        self,
        data_source: DataSource,
        fail_on_error: bool = True,
        logging_context: str = "",
    ) -> None:
        """Add data source semantic information, validating consistency with existing data sources."""
        errors = []

        if data_source.name in self._data_source_names:
            errors.append(f"name {data_source.name} already registered - please ensure data source names are unique")

        for measure in data_source.measures:
            if measure.reference in self._measure_aggs and self._measure_aggs[measure.reference] != measure.agg:
                errors.append(
                    f"conflicting aggregation (agg) for measure `{measure.reference.element_name}` registered as "
                    f"`{self._measure_aggs[measure.reference]}`; Got `{measure.agg}"
                )

        if errors:
            error_prefix = "\n  - "
            error_msg = (
                f"Unable to add data source `{data_source.name}` "
                f"{'while ' + logging_context + ' ' if logging_context else ''}"
                f"{'... skipping' if not fail_on_error else ''}.\n"
                f"Errors: {error_prefix + error_prefix.join(errors)}"
            )
            if fail_on_error:
                raise InvalidDataSourceError(error_msg)
            logger.warning(error_msg)
            return

        self._data_source_names.add(data_source.name)
        self._data_source_to_aggregation_time_dimensions[data_source.reference] = ElementGrouper[
            TimeDimensionReference, MeasureSpec
        ]()

        for measure in data_source.measures:
            self._measure_aggs[measure.reference] = measure.agg
            self._measure_index[measure.reference].append(data_source)
            agg_time_dimension = measure.checked_agg_time_dimension
            self._data_source_to_aggregation_time_dimensions[data_source.reference].add_value(
                key=agg_time_dimension,
                value=MeasureConverter.convert_to_measure_spec(measure=measure),
            )
            self._measure_agg_time_dimension[measure.reference] = agg_time_dimension
            if measure.non_additive_dimension:
                non_additive_dimension_spec = NonAdditiveDimensionSpec(
                    name=measure.non_additive_dimension.name,
                    window_choice=measure.non_additive_dimension.window_choice,
                    window_groupings=tuple(measure.non_additive_dimension.window_groupings),
                )
                self._measure_non_additive_dimension_specs[measure.reference] = non_additive_dimension_spec
        for dim in data_source.dimensions:
            self._linkable_reference_index[dim.reference].append(data_source)
            self._dimension_index[dim.reference].append(data_source)
        for ident in data_source.identifiers:
            self._identifier_ref_to_entity[ident.reference] = ident.entity
            self._entity_index[ident.entity].append(data_source)
            self._linkable_reference_index[ident.reference].append(data_source)

    @property
    def data_source_references(self) -> Sequence[DataSourceReference]:  # noqa: D
        data_source_names_sorted = sorted(self._data_source_names)
        return tuple(DataSourceReference(data_source_name=x) for x in data_source_names_sorted)

    def get_aggregation_time_dimensions_with_measures(
        self, data_source_reference: DataSourceReference
    ) -> ElementGrouper[TimeDimensionReference, MeasureSpec]:
        """Return all time dimensions in a data source with their associated measures."""
        assert (
            data_source_reference in self._data_source_to_aggregation_time_dimensions
        ), f"Data Source {data_source_reference} is not known"
        return self._data_source_to_aggregation_time_dimensions[data_source_reference]

    def get_data_sources_for_identifier(self, identifier_reference: IdentifierReference) -> Set[DataSource]:
        """Return all data sources associated with an identifier reference"""
        identifier_entity = self._identifier_ref_to_entity[identifier_reference]
        return set(self._entity_index[identifier_entity])
