"""
Each of these Interface subclasses provides a protocol on how
the corresponding Adapter subclass should be implemented.

Usage: TODO: add usage info

"""

from __future__ import annotations
import importlib

import logging

from abc import abstractmethod
from collections import defaultdict
from peh_model import peh
from typing import TYPE_CHECKING, Callable, Generic, cast, List

from pypeh.core.cache.containers import CacheContainerView
from pypeh.core.models.constants import ObservablePropertyValueType
from pypeh.core.models.internal_data_layout import Dataset, DatasetSchemaElement, DatasetSeries
from pypeh.core.models.typing import T_DataType
from pypeh.core.models.settings import FileSystemSettings
from pypeh.core.models import graph, validation_dto
from pypeh.core.session.connections import ConnectionManager

if TYPE_CHECKING:
    from typing import Sequence, Any
    from pypeh.core.models.validation_errors import ValidationErrorReport

logger = logging.getLogger(__name__)


class OutDataOpsInterface(Generic[T_DataType]):
    """
    Example of DataOps methods
    def validate(self, data: Mapping, config: Mapping):
        pass

    def summarize(self, dat: Mapping, config: Mapping):
        pass
    """

    @classmethod
    def get_default_adapter_class(cls):
        try:
            adapter_module = importlib.import_module("pypeh.adapters.outbound.dataops.dataframe_adapter")
            adapter_class = getattr(adapter_module, "DataFrameAdapter")
        except Exception as e:
            logger.error("Exception encountered while attempting to import dataops DataFrameAdapter")
            raise e
        return adapter_class

    @abstractmethod
    def _join_data(
        self,
        data: T_DataType,
        other_data: list[T_DataType],
        join_on: list[str],
        subset_fields_other: list[list[str]],
    ) -> T_DataType:
        raise NotImplementedError

    @abstractmethod
    def select_field(self, dataset, field_label: str):
        raise NotImplementedError

    @abstractmethod
    def get_element_labels(self, data: T_DataType) -> list[str]:
        raise NotImplementedError

    @abstractmethod
    def get_element_values(self, data: T_DataType, element_label: str, as_list=True) -> set[str] | list[str]:
        raise NotImplementedError

    @abstractmethod
    def check_element_has_empty_values(self, data: T_DataType, element_label: str) -> bool:
        raise NotImplementedError

    @abstractmethod
    def check_element_has_only_empty_values(self, data: T_DataType, element_label: str) -> bool:
        raise NotImplementedError

    @abstractmethod
    def subset(
        self,
        data: T_DataType,
        element_group: list[str],
        id_group: list[tuple[Any]] | None = None,
        identifying_elements: list[str] | None = None,
    ) -> T_DataType: ...

    def relabel(self, data: T_DataType, element_mapping: dict[str, str]) -> T_DataType: ...

    @abstractmethod
    def collect(self, datasets: dict):
        raise NotImplementedError

    @abstractmethod
    def type_mapper(self, peh_value_type: str | ObservablePropertyValueType):
        raise NotImplementedError


class ValidationInterface(OutDataOpsInterface, Generic[T_DataType]):
    @abstractmethod
    def _validate(
        self, data: dict[str, Sequence] | T_DataType, config: validation_dto.ValidationConfig
    ) -> ValidationErrorReport:
        raise NotImplementedError

    @classmethod
    def get_default_adapter_class(cls):
        try:
            adapter_module = importlib.import_module(
                "pypeh.adapters.outbound.validation.pandera_adapter.validation_adapter"
            )
            adapter_class = getattr(adapter_module, "DataFrameValidationAdapter")
        except Exception as e:
            logger.error("Exception encountered while attempting to import a Pandera-based DataFrameAdapter")
            raise e
        return adapter_class

    def build_column_validation(
        self,
        dataset_schema_element: DatasetSchemaElement,
        type_annotations: dict[str, dict[str, ObservablePropertyValueType]],
        cache_view: CacheContainerView,
        dataset_label: str | None = None,
        column_has_only_empty_values: bool = False,
        allow_incomplete: bool = False,
    ) -> validation_dto.ColumnValidation:
        # Set default validation check applicability flags
        apply_required_check = True
        apply_nullable_check = True
        apply_property_validation = True
        # Adapt the flags in accordance with the user request, project config and data
        if allow_incomplete:
            apply_required_check = False
            if column_has_only_empty_values:
                apply_nullable_check = False
                apply_property_validation = False

        validations = []
        observable_property_id = dataset_schema_element.observable_property_id
        observable_property = cache_view.get(observable_property_id, "ObservableProperty")
        assert isinstance(
            observable_property, peh.ObservableProperty
        ), f"ObservableProperty with id {observable_property_id} not found"

        if apply_required_check:
            required = observable_property.default_required
        else:
            required = False

        if apply_nullable_check:
            nullable = not required  # required and nullable are now checking the same thing
        else:
            nullable = True

        if apply_property_validation:
            if validation_designs := getattr(observable_property, "validation_designs", None):
                validations.extend(
                    [
                        validation_dto.ValidationDesign.from_peh(
                            vd, type_annotations=type_annotations, dataset_label=dataset_label
                        )
                        for vd in validation_designs
                    ]
                )
            if value_metadata := getattr(observable_property, "value_metadata", None):
                validations.extend(
                    validation_dto.ValidationDesign.list_from_metadata(
                        value_metadata, type_annotations=type_annotations, dataset_label=dataset_label
                    )
                )
            if getattr(observable_property, "categorical", None):
                value_options = getattr(observable_property, "value_options", None)
                assert (
                    value_options is not None
                ), f"ObservableProperty {observable_property} lacks `value_options` for categorical type"
                assert dataset_schema_element.data_type == ObservablePropertyValueType.STRING
                validation_arg_values: list[str] = [str(vo.key) for vo in value_options]
                # TODO: ADD CUSTOM CHECKS ON CATEGORICAL VARIABLES HERE
                expr = validation_dto.ValidationExpression(
                    command="is_in",
                    arg_values=validation_arg_values,
                    arg_columns=None,
                    subject=None,
                )
                validation = validation_dto.ValidationDesign(
                    name="check_categorical",
                    error_level=validation_dto.ValidationErrorLevel.ERROR,
                    expression=expr,
                    error_message=None,  # CUSTOM ERROR MESSAGE CANNOT BE PROVIDED OR SHOULD BE HARDCODED
                )
                validations.append(validation)

            if dataset_schema_element.data_type in [
                ObservablePropertyValueType.CATEGORICAL,
                ObservablePropertyValueType.STRING,
            ]:
                expr = validation_dto.ValidationExpression(
                    command="trailing_spaces",
                    arg_values=None,
                    arg_columns=None,
                    subject=None,
                )
                validation = validation_dto.ValidationDesign(
                    name="check_trailing_spaces",
                    error_level=validation_dto.ValidationErrorLevel.ERROR,
                    expression=expr,
                    error_message="Trailing/leading spaces were detected. Remove unnecessary white spaces.",
                )
                validations.append(validation)

        assert dataset_schema_element.data_type.value != "decimal"
        # transformation using context_magic
        assert isinstance(required, bool)
        return validation_dto.ColumnValidation(
            unique_name=dataset_schema_element.label,
            data_type=dataset_schema_element.data_type.value,
            required=required,
            nullable=nullable,
            validations=validations,
        )

    def collect_column_validations(
        self,
        dataset: Dataset,
        type_annotations: dict[str, dict[str, ObservablePropertyValueType]],
        cache_view: CacheContainerView,
        allow_incomplete: bool = False,
    ) -> list[validation_dto.ColumnValidation]:
        column_validations: list[validation_dto.ColumnValidation] = []
        column_labels = dataset.get_element_labels()
        assert column_labels is not None
        for column_label in column_labels:
            dataset_schema_element = dataset.get_schema_element_by_label(column_label)
            assert dataset_schema_element is not None
            # Check whether the dataset has data in the column
            if dataset.data is None:
                column_has_only_empty_values = True
            else:
                column_has_only_empty_values = self.check_element_has_only_empty_values(
                    data=dataset.data,
                    element_label=column_label,
                )
            column_validation = self.build_column_validation(
                dataset_schema_element=dataset_schema_element,
                cache_view=cache_view,
                type_annotations=type_annotations,
                column_has_only_empty_values=column_has_only_empty_values,
                allow_incomplete=allow_incomplete,
            )
            column_validations.append(column_validation)

        return column_validations

    def build_dataset_level_validations(
        self,
        dataset: Dataset,
        dataset_series: DatasetSeries | None = None,
        cache_view: CacheContainerView | None = None,
    ) -> list[validation_dto.ValidationDesign] | None:
        dataset_level_validations: list[validation_dto.ValidationDesign] = []
        if cache_view is None:
            raise NotImplementedError

        if dataset_series is not None:
            type_annotations = dataset_series.get_type_annotations()
        else:
            type_annotations = {dataset.label: dataset.get_type_annotations()}
        layout_section_id = dataset.described_by
        if layout_section_id is None:
            return None
        assert isinstance(layout_section_id, str)
        layout_section = cache_view.get(layout_section_id, "DataLayoutSection")
        assert layout_section is not None
        assert isinstance(layout_section, peh.DataLayoutSection)
        if layout_section.validation_designs:
            for vd in layout_section.validation_designs:
                assert isinstance(vd, peh.ValidationDesign)
                dataset_validation = validation_dto.ValidationDesign.from_peh(vd, type_annotations)
                # For an expression that relies on a field reference spec for its arguments, set the validation arguments
                # as the actual values from the dataset (e.g. for an "is_in" check on a foreign key relation)
                validation_expression = vd.validation_expression
                assert isinstance(validation_expression, peh.ValidationExpression)
                if validation_expression.validation_arg_contextual_field_references:
                    arg_values = validation_expression.validation_arg_values
                    assert isinstance(arg_values, list)
                    for ref in validation_expression.validation_arg_contextual_field_references:
                        assert isinstance(ref, peh.ContextualFieldReference)
                        dataset_label = ref.dataset_label
                        assert dataset_label is not None
                        field_label = ref.field_label
                        assert field_label is not None
                        if dataset_series is None:
                            assert dataset_label == dataset.label
                            dependent_dataset = dataset
                        else:
                            dependent_dataset = dataset_series[dataset_label]
                        assert dependent_dataset is not None
                        # fix this at the Interface level
                        column_arg_values = self.get_element_values(dependent_dataset.data, field_label, as_list=True)
                        arg_values.extend(column_arg_values)
                    dataset_validation.expression.arg_values = arg_values
                    dataset_validation.expression.arg_columns = None
                    dataset_level_validations.append(dataset_validation)

        return dataset_level_validations

    @classmethod
    def merge_contextual_field_reference_dependencies(
        cls, column_validations: list[validation_dto.ColumnValidation]
    ) -> dict[str, set[str]]:
        dependent_contextual_field_references = defaultdict(set)
        for column_validation in column_validations:
            if column_validation.validations is not None:
                for validation_design in column_validation.validations:
                    dependency_dict = validation_design.dependent_contextual_field_references
                    if dependency_dict is not None:
                        for ds, fields in dependency_dict.items():
                            dependent_contextual_field_references[ds].update(fields)

        return dependent_contextual_field_references

    def build_validation_config(
        self,
        dataset: Dataset,
        dataset_series: DatasetSeries | None = None,
        cache_view: CacheContainerView | None = None,
        allow_incomplete: bool = False,
    ) -> validation_dto.ValidationConfig:
        column_validations = []
        dataset_validations = []
        if cache_view is None:
            raise NotImplementedError("The absence of a CacheView is currently not supported")

        identifying_column_names = dataset.get_primary_keys()
        if identifying_column_names is not None:
            identifying_column_names = list(identifying_column_names)

        if dataset_series is None:
            type_annotations = {dataset.label: dataset.get_type_annotations()}
        else:
            type_annotations = dataset_series.get_type_annotations()
        column_validations = self.collect_column_validations(
            dataset=dataset,
            type_annotations=type_annotations,
            cache_view=cache_view,
            allow_incomplete=allow_incomplete,
        )
        dependent_contextual_field_references = self.merge_contextual_field_reference_dependencies(column_validations)
        dataset_validations = self.build_dataset_level_validations(
            dataset=dataset,
            dataset_series=dataset_series,
            cache_view=cache_view,
        )

        return validation_dto.ValidationConfig(
            name=dataset.label,
            columns=column_validations,
            identifying_column_names=identifying_column_names,
            validations=dataset_validations,
            dependent_contextual_field_references=dependent_contextual_field_references,
        )

    def validate(
        self,
        dataset: Dataset[T_DataType],
        dependent_dataset_series: DatasetSeries[T_DataType] | None = None,
        cache_view: CacheContainerView | None = None,
        allow_incomplete: bool = False,
    ) -> ValidationErrorReport:
        assert dataset.data is not None
        assert cache_view is not None
        to_validate = dataset.data
        assert to_validate is not None

        validation_config = self.build_validation_config(
            dataset=dataset,
            dataset_series=dependent_dataset_series,
            cache_view=cache_view,
            allow_incomplete=allow_incomplete,
        )
        # check whether data requires join to perform validation (cross DataLayoutSection validation)
        join_required = False
        dependent_contextual_field_references = validation_config.dependent_contextual_field_references
        if dependent_contextual_field_references is not None:
            join_required = len(dependent_contextual_field_references) > 0

        if join_required:
            if dependent_dataset_series is None:
                me = "`dependent_data` is required to perform all validations with `ValidationInterface`"
                logger.error(me)
                raise ValueError(me)
            else:
                assert (
                    dependent_contextual_field_references is not None
                ), "dependent_contextual_field_references in `ValidationInterface.validate` should not be None"
                assert dependent_dataset_series is not None
                identifying_field_labels = dataset.schema.primary_keys
                assert identifying_field_labels is not None
                all_other_data = []
                subset_fields_other = []

                for dataset_label, dependent_field_labels in dependent_contextual_field_references.items():
                    other_dataset = dependent_dataset_series[dataset_label]
                    assert other_dataset is not None
                    other_data = other_dataset.data
                    assert other_data is not None
                    all_other_data.append(other_data)
                    field_subset = [*identifying_field_labels, *dependent_field_labels]
                    subset_fields_other.append(field_subset)
                    to_validate = self._join_data(
                        data=to_validate,
                        other_data=all_other_data,
                        join_on=list(identifying_field_labels),
                        subset_fields_other=subset_fields_other,
                    )

        ret = self._validate(to_validate, validation_config)

        return ret


class DataEnrichmentInterface(OutDataOpsInterface, Generic[T_DataType]):
    @classmethod
    def get_default_adapter_class(cls):
        try:
            adapter_module = importlib.import_module("pypeh.adapters.outbound.enrichment.dataframe_adapter")
            adapter_class = getattr(adapter_module, "DataFrameEnrichmentAdapter")
        except Exception as e:
            logger.error("Exception encountered while attempting to import enrichment DataFrameAdapter")
            raise e
        return adapter_class

    @abstractmethod
    def apply_joins(self, datasets, join_specs, node): ...

    @abstractmethod
    def apply_map(self, dataset, map_fn, field_label, output_dtype, base_fields, **kwargs): ...

    @abstractmethod
    def map_type(self, peh_value_type: str): ...

    def build_callable(self, delayed_node: graph.Delayed) -> Callable:
        map_fn = delayed_node.map_fn
        arg_sources = delayed_node.arg_sources
        join_specs = delayed_node.join_specs
        output_dtype = self.map_type(delayed_node.output_dtype)

        def _apply(datasets: dict, *, node: graph.Node, base_fields: dict):
            """
            datasets: dict of dataset_label → dataset object (lazy or eager)
            parent_results: mapping from parent Node → computed result for this node
            """
            ds = datasets[node.dataset_label]
            # Apply all joins
            if join_specs:
                ds = self.apply_joins(
                    datasets=datasets,
                    join_specs=join_specs,
                    node=node,
                )

            # Build column expressions for the map function
            kwargs = {}
            for arg_name, parent_node in arg_sources.items():
                col_name = parent_node.field_label
                kwargs[arg_name] = self.select_field(ds, col_name)

            # Apply the map
            base_fields_subset = base_fields.get(node.dataset_label, None)
            assert base_fields_subset is not None
            out = self.apply_map(ds, map_fn, node.field_label, output_dtype, base_fields_subset, **kwargs)
            base_fields_subset.append(node.field_label)

            return out

        return _apply

    def compile_dependency_graph(self, dependency_graph: graph.Graph) -> graph.ExecutionPlan:
        sorted_nodes = dependency_graph.topological_sort()
        steps: list[graph.ExecutionStep] = []

        for node in sorted_nodes:
            delayed = dependency_graph.delayed_fns.get(node)
            if delayed is None:
                continue

            compute_fn = self.build_callable(delayed)
            steps.append(graph.ExecutionStep(node=node, compute=compute_fn))

        ret = graph.ExecutionPlan(steps)
        dependency_graph.execution_plan = ret

        return ret

    def compute_with_dependency_graph(self, dependency_graph: graph.Graph, datasets: dict[str, Dataset]):
        if dependency_graph.execution_plan is None:
            raise AssertionError("A dependency graph needs to be compiled first to set up an execution plan")

        raw_datasets = {label: dataset.data for label, dataset in datasets.items()}
        base_fields = {label: dataset.get_element_labels() for label, dataset in datasets.items()}

        dependency_graph.execution_plan.run(raw_datasets, base_fields)
        self.collect(datasets)
        for dataset_label in datasets:
            datasets[dataset_label].data = raw_datasets[dataset_label]

    def build_dependency_graph(
        self,
        observations: list[peh.Observation],
        contextual_field_reference_map: dict[str, tuple[str, str] | None],
        cache_view: CacheContainerView,
        join_spec_mapping: dict | None = None,
    ) -> graph.Graph:
        dependency_graph = graph.Graph()
        nested_entity_paths = [
            ["observation_design", "identifying_observable_property_id_list"],
            ["observation_design", "required_observable_property_id_list"],
            ["observation_design", "optional_observable_property_id_list"],
        ]
        # the source_observations also need to be added to the dependency graph!!!!!
        for observation in observations:
            observation_id = observation.id
            for path in nested_entity_paths:
                for observable_property in cache_view.walk_entity(
                    entity_id=observation_id, nested_entity_path=path, entity_type="Observation"
                ):
                    assert isinstance(observable_property, peh.ObservableProperty)
                    target_contextual_field_ref = contextual_field_reference_map.get(observable_property.id)
                    assert (
                        target_contextual_field_ref is not None
                    ), f"Observable property with id {observable_property.id} is non unique."
                    target_dataset_label, target_field_label = target_contextual_field_ref
                    calculation_designs = observable_property.calculation_designs
                    if calculation_designs:
                        # EXTRA INFO FROM CALCULATION DESIGN AND UPDATE DEPENDENCY GRAPH
                        child = graph.Node(dataset_label=target_dataset_label, field_label=target_field_label)
                        if len(calculation_designs) > 1:
                            raise NotImplementedError(
                                "No current implementation for multiple calculation designs linked to a single observable property"
                            )
                        calculation_design = calculation_designs[0]
                        assert isinstance(calculation_design, peh.CalculationDesign)
                        calculation_implementation = calculation_design.calculation_implementation
                        assert isinstance(calculation_implementation, peh.CalculationImplementation)
                        output_dtype = observable_property.value_type
                        assert output_dtype is not None
                        function_name = calculation_implementation.function_name
                        assert isinstance(function_name, str)

                        dependency_graph.add_calculation_target(
                            child, function_name=function_name, result_dtype=output_dtype
                        )
                        function_kwargs = calculation_implementation.function_kwargs
                        assert function_kwargs is not None
                        for function_kwarg in function_kwargs:
                            assert isinstance(function_kwarg, peh.CalculationKeywordArgument)
                            source_field_ref = function_kwarg.contextual_field_reference
                            assert isinstance(source_field_ref, peh.ContextualFieldReference)
                            source_dataset_label = source_field_ref.dataset_label
                            assert source_dataset_label is not None
                            source_field_label = source_field_ref.field_label
                            assert source_field_label is not None
                            parent = graph.Node(dataset_label=source_dataset_label, field_label=source_field_label)
                            map_name = function_kwarg.mapping_name
                            assert map_name is not None
                            join_spec = None
                            if join_spec_mapping is not None:
                                join_spec = join_spec_mapping.get(
                                    frozenset([target_dataset_label, source_dataset_label]), None
                                )
                                if join_spec is not None:
                                    assert len(join_spec) == 1, "Complex JoinSpecs are not supported yet."
                            dependency_graph.add_calculation_source(
                                parent,
                                child,
                                map_name,
                                join_spec=join_spec,
                            )

        return dependency_graph

    def enrich(
        self,
        source_dataset_series: DatasetSeries,
        target_observations: list[peh.Observation],
        target_derived_from: list[peh.Observation],
        cache_view: CacheContainerView,
    ) -> DatasetSeries:
        # ADD TARGET OBSERVATION TO SOURCE_DATASET_SERIES
        for source_obs, target_observation in zip(target_derived_from, target_observations):
            source_dataset = source_dataset_series.get_dataset_by_observation(source_obs.id, rebuild_index=True)
            assert source_dataset is not None
            source_dataset_label = source_dataset.label
            source_dataset_series.add_observation(
                source_dataset_label,
                target_observation,
                cache_view,
            )
        join_spec_mapping = source_dataset_series.resolve_all_joins()
        # BUILD DEPENDENCY GRAPH
        all_observations = []
        for nested_observations in source_dataset_series.observations.values():
            for observation_id in nested_observations:
                observation = cache_view.get(observation_id, "Observation")
                assert observation is not None
                all_observations.append(observation)

        dependency_graph = self.build_dependency_graph(
            observations=all_observations,
            contextual_field_reference_map=source_dataset_series.get_contextual_field_reference_index(),
            join_spec_mapping=join_spec_mapping,
            cache_view=cache_view,
        )
        # EXECUTE THE DEFINED COMPUTATIONS
        self.compile_dependency_graph(dependency_graph=dependency_graph)
        self.compute_with_dependency_graph(
            dependency_graph=dependency_graph,
            datasets=source_dataset_series.parts,
        )
        # RETURN THE UPDATED SOURCE_DATASET_SERIES
        return source_dataset_series


class DataImportInterface(OutDataOpsInterface, Generic[T_DataType]):
    @classmethod
    def get_default_adapter_class(cls):
        try:
            adapter_module = importlib.import_module(
                "pypeh.adapters.outbound.validation.pandera_adapter.validation_adapter"
            )
            adapter_class = getattr(adapter_module, "DataFrameValidationAdapter")
        except Exception as e:
            logger.error("Exception encountered while attempting to import a Pandera-based DataFrameAdapter")
            raise e
        return adapter_class

    @abstractmethod
    def import_data(self, source: str, config: FileSystemSettings) -> T_DataType | List[T_DataType]:
        raise NotImplementedError

    def import_data_layout(
        self,
        source: str,
        config: FileSystemSettings,
        **kwargs,
    ) -> peh.DataLayout | List[peh.DataLayout]:
        provider = ConnectionManager._create_adapter(config)
        layout = provider.load(source)
        if isinstance(layout, peh.EntityList):
            layout = layout.layouts
        if isinstance(layout, list):
            if not all(isinstance(item, peh.DataLayout) for item in layout):
                me = "Imported layouts should all be DataLayout instances"
                logger.error(me)
                raise TypeError(me)
            return cast(List[peh.DataLayout], layout)

        elif isinstance(layout, peh.DataLayout):
            return layout

        else:
            me = f"Imported layout should be a DataLayout instance not {type(layout)}"
            logger.error(me)
            raise TypeError(me)

    def _extract_observed_value_provenance(self) -> bool:
        return True

    def _normalize_observable_properties(self) -> bool:
        raise NotImplementedError

    def _raw_data_to_observation_data(
        self,
        raw_data: T_DataType,
        data_layout_element_labels: list[str],
        identifying_layout_element_label: str,
        entity_id_list: list[str] | None = None,
    ) -> T_DataType:
        raise NotImplementedError


class AggregationInterface(OutDataOpsInterface, Generic[T_DataType]):
    @abstractmethod
    def _summarize(
        self, df: T_DataType, group_cols: list[str], value_col: str, stat_builders: list, **kwargs
    ) -> T_DataType:
        raise NotImplementedError

    @classmethod
    def get_default_adapter_class(cls):
        try:
            adapter_module = importlib.import_module(
                "pypeh.adapters.outbound.aggregation.polars_adapter.dataframe_adapter"
            )

            adapter_class = getattr(adapter_module, "DataFrameAggregationAdapter")
        except Exception as e:
            logger.error("Exception encountered while attempting to import a Polars-based DataFrameAggregationAdapter")
            raise e
        return adapter_class
