from abc import ABC
from datetime import datetime
from enum import Enum
from typing import Any, Optional, Tuple, Union

from pydantic import BaseModel, Extra
from typing_extensions import Literal

from acryl_datahub_cloud.sdk.assertion_input.assertion_input import (
    DEFAULT_EVERY_SIX_HOURS_SCHEDULE,
    AssertionIncidentBehavior,
    DetectionMechanismInputTypes,
    FieldSpecType,
    _AssertionInput,
)
from acryl_datahub_cloud.sdk.entities.assertion import (
    Assertion,
    AssertionInfoInputType,
    TagsInputType,
)
from acryl_datahub_cloud.sdk.errors import SDKNotYetSupportedError, SDKUsageError
from datahub.metadata import schema_classes as models
from datahub.metadata.urns import (
    AssertionUrn,
    CorpUserUrn,
    DatasetUrn,
)
from datahub.sdk.entity_client import EntityClient

# TODO: better naming for "volume assertion definition"


class VolumeAssertionOperator(str, Enum):
    """Valid operators for volume assertions."""

    LESS_THAN_OR_EQUAL_TO = models.AssertionStdOperatorClass.LESS_THAN_OR_EQUAL_TO
    GREATER_THAN_OR_EQUAL_TO = models.AssertionStdOperatorClass.GREATER_THAN_OR_EQUAL_TO
    BETWEEN = models.AssertionStdOperatorClass.BETWEEN


class _AbstractVolumeAssertionDefinition(BaseModel, ABC):
    type: str

    class Config:
        extra = Extra.forbid


class RowCountTotal(_AbstractVolumeAssertionDefinition):
    type: Literal["row_count_total"] = "row_count_total"
    operator: VolumeAssertionOperator
    parameters: Union[float, Tuple[float, float]]


class RowCountChange(_AbstractVolumeAssertionDefinition):
    type: Literal["row_count_change"] = "row_count_change"
    kind: Literal["absolute", "percent"]
    operator: VolumeAssertionOperator
    parameters: Union[float, Tuple[float, float]]


_VOLUME_ASSERTION_DEFINITION_CONCRETE_TYPES = (
    RowCountTotal,
    RowCountChange,
)
_VolumeAssertionDefinitionTypes = Union[
    RowCountTotal,
    RowCountChange,
]

VolumeAssertionDefinitionInputTypes = Union[
    dict[str, Any], _VolumeAssertionDefinitionTypes
]


class VolumeAssertionDefinition:
    ROW_COUNT_TOTAL = RowCountTotal
    ROW_COUNT_CHANGE = RowCountChange

    @staticmethod
    def _validate_between_parameters(
        parameters: Union[float, Tuple[float, float]], assertion_type: str
    ) -> None:
        """Validate parameters for BETWEEN operator."""
        if not isinstance(parameters, tuple) or len(parameters) != 2:
            raise SDKUsageError(
                f"For BETWEEN operator in {assertion_type}, parameters must be a tuple of two numbers (min_value, max_value)."
            )

    @staticmethod
    def _validate_single_value_parameters(
        parameters: Union[float, Tuple[float, float]],
        operator_enum: VolumeAssertionOperator,
        assertion_type: str,
    ) -> None:
        """Validate parameters for single-value operators."""
        if not isinstance(parameters, (int, float)):
            if isinstance(parameters, tuple):
                raise SDKUsageError(
                    f"For {operator_enum.value} operator in {assertion_type}, parameters must be a single number, not a tuple."
                )
            else:
                raise SDKUsageError(
                    f"For {operator_enum.value} operator in {assertion_type}, parameters must be a single number."
                )

    @staticmethod
    def _parse_operator(
        operator: Union[str, VolumeAssertionOperator],
    ) -> VolumeAssertionOperator:
        """Parse and validate operator input, converting string to enum if needed."""
        if isinstance(operator, str):
            try:
                return VolumeAssertionOperator(operator)
            except ValueError as e:
                valid_operators = ", ".join(
                    [op.value for op in VolumeAssertionOperator]
                )
                raise SDKUsageError(
                    f"Invalid operator '{operator}'. Valid operators: {valid_operators}"
                ) from e
        return operator

    @staticmethod
    def _validate_operator_and_parameters(
        operator: Union[str, VolumeAssertionOperator],
        parameters: Union[float, Tuple[float, float]],
        assertion_type: str,
    ) -> None:
        """Validate that operator and parameters are compatible for volume assertions."""
        operator_enum = VolumeAssertionDefinition._parse_operator(operator)

        # Validate parameter structure based on operator
        if operator_enum == VolumeAssertionOperator.BETWEEN:
            VolumeAssertionDefinition._validate_between_parameters(
                parameters, assertion_type
            )
        else:
            VolumeAssertionDefinition._validate_single_value_parameters(
                parameters, operator_enum, assertion_type
            )

    @staticmethod
    def _parse_instantiated_object(
        definition: _VolumeAssertionDefinitionTypes,
    ) -> _VolumeAssertionDefinitionTypes:
        """Parse and validate already instantiated volume assertion objects."""
        VolumeAssertionDefinition._validate_operator_and_parameters(
            definition.operator, definition.parameters, definition.type
        )
        return definition

    @staticmethod
    def _parse_dict_definition(
        definition_dict: dict[str, Any],
    ) -> _VolumeAssertionDefinitionTypes:
        """Parse and validate dictionary-based volume assertion definitions."""
        try:
            assertion_type = definition_dict.pop("type")
        except KeyError as e:
            raise SDKUsageError(
                "Volume assertion definition must include a 'type' field"
            ) from e

        # Check for valid assertion type first
        if assertion_type not in ["row_count_total", "row_count_change"]:
            raise SDKUsageError(
                f"Unknown volume assertion type: {assertion_type}. Supported types: row_count_total, row_count_change"
            )

        # Extract operator and parameters for validation
        operator = definition_dict.get("operator")
        parameters = definition_dict.get("parameters")

        if operator is None:
            raise SDKUsageError(
                f"Missing required 'operator' field for {assertion_type}"
            )
        if parameters is None:
            raise SDKUsageError(
                f"Missing required 'parameters' field for {assertion_type}"
            )

        # Validate basic parameter type first
        if not isinstance(parameters, (int, float, tuple)):
            raise SDKUsageError(
                f"For {assertion_type}, parameters must be a number or a tuple of two numbers, got: {type(parameters)}"
            )

        # Validate operator and parameters before object creation
        VolumeAssertionDefinition._validate_operator_and_parameters(
            operator, parameters, assertion_type
        )

        # Convert string operator to enum for object creation
        if isinstance(operator, str):
            definition_dict["operator"] = VolumeAssertionDefinition._parse_operator(
                operator
            )

        if assertion_type == "row_count_total":
            try:
                return RowCountTotal(**definition_dict)
            except Exception as e:
                raise SDKUsageError(
                    f"Failed to create row_count_total volume assertion: {str(e)}"
                ) from e
        else:  # assertion_type == "row_count_change"
            try:
                return RowCountChange(**definition_dict)
            except Exception as e:
                raise SDKUsageError(
                    f"Failed to create row_count_change volume assertion: {str(e)}"
                ) from e

    @staticmethod
    def parse(
        definition: VolumeAssertionDefinitionInputTypes,
    ) -> _VolumeAssertionDefinitionTypes:
        """Parse and validate a volume assertion definition.

        This method converts dictionary-based volume assertion definitions into typed volume
        assertion objects, or validates already instantiated volume assertion objects. It
        supports two volume assertion types: row_count_total and row_count_change.

        Args:
            definition: A volume assertion definition that can be either:
                - A dictionary containing volume assertion configuration with keys:
                  - type: Must be "row_count_total" or "row_count_change"
                  - operator: Must be "LESS_THAN_OR_EQUAL_TO", "GREATER_THAN_OR_EQUAL_TO", or "BETWEEN"
                  - parameters: Number for single-value operators, tuple of two numbers for BETWEEN
                  - kind: Required for "row_count_change", must be "absolute" or "percent"
                - An already instantiated RowCountTotal or RowCountChange object

        Returns:
            A validated volume assertion definition object (RowCountTotal or RowCountChange).

        Raises:
            SDKUsageError: If the definition is invalid, including:
                - Invalid input type (not dict or volume assertion object)
                - Missing required fields (type, operator, parameters, kind for row_count_change)
                - Unknown assertion type (not row_count_total or row_count_change)
                - Invalid operator (not in allowed operators)
                - Invalid parameter structure for operator:
                  - Single-value operators require number parameters
                  - BETWEEN operator requires tuple of two numbers
                - Object construction failures (extra fields, validation errors)

        Examples:
            Parse a row count total assertion:
            >>> definition = {
            ...     "type": "row_count_total",
            ...     "operator": "GREATER_THAN_OR_EQUAL_TO",
            ...     "parameters": 100
            ... }
            >>> result = VolumeAssertionDefinition.parse(definition)
            >>> isinstance(result, RowCountTotal)
            True

            Parse a row count change assertion with BETWEEN operator:
            >>> definition = {
            ...     "type": "row_count_change",
            ...     "kind": "absolute",
            ...     "operator": "BETWEEN",
            ...     "parameters": (10, 50)
            ... }
            >>> result = VolumeAssertionDefinition.parse(definition)
            >>> isinstance(result, RowCountChange)
            True

            Parse an already instantiated object:
            >>> obj = RowCountTotal(
            ...     operator=VolumeAssertionOperator.LESS_THAN_OR_EQUAL_TO,
            ...     parameters=200
            ... )
            >>> result = VolumeAssertionDefinition.parse(obj)
            >>> result == obj
            True
        """
        # If already instantiated, validate and return
        if isinstance(definition, _VOLUME_ASSERTION_DEFINITION_CONCRETE_TYPES):
            return VolumeAssertionDefinition._parse_instantiated_object(definition)

        if not isinstance(definition, dict):
            raise SDKUsageError(
                f"Volume assertion definition must be a dict or a volume assertion definition object, got: {type(definition)}"
            )

        return VolumeAssertionDefinition._parse_dict_definition(definition.copy())

    @staticmethod
    def build_model_volume_info(
        definition: _VolumeAssertionDefinitionTypes,
        dataset_urn: str,
        filter: Optional[models.DatasetFilterClass] = None,
    ) -> models.VolumeAssertionInfoClass:
        """Build a DataHub VolumeAssertionInfoClass from a validated volume assertion definition.

        This method converts validated volume assertion definition objects into DataHub model
        classes suitable for creating volume assertions in the DataHub metadata service.

        Args:
            definition: A validated volume assertion definition object (RowCountTotal or RowCountChange).
                        This should be the output of VolumeAssertionDefinition.parse().
            dataset_urn: The dataset URN that this assertion applies to.
            filter: Optional filter to apply to the assertion.

        Returns:
            A VolumeAssertionInfoClass configured for the specific volume assertion type.

        Raises:
            SDKUsageError: If the definition type is not supported.
        """
        if isinstance(definition, RowCountTotal):
            volume_info = models.VolumeAssertionInfoClass(
                type=models.VolumeAssertionTypeClass.ROW_COUNT_TOTAL,
                entity=dataset_urn,
                rowCountTotal=models.RowCountTotalClass(
                    operator=definition.operator.value,
                    parameters=VolumeAssertionDefinition._build_assertion_parameters(
                        definition.operator, definition.parameters
                    ),
                ),
            )
            if filter is not None:
                volume_info.filter = filter
            return volume_info
        elif isinstance(definition, RowCountChange):
            # Map kind to DataHub assertion value change type
            change_type = (
                models.AssertionValueChangeTypeClass.ABSOLUTE
                if definition.kind == "absolute"
                else models.AssertionValueChangeTypeClass.PERCENTAGE
            )

            volume_info = models.VolumeAssertionInfoClass(
                type=models.VolumeAssertionTypeClass.ROW_COUNT_CHANGE,
                entity=dataset_urn,
                rowCountChange=models.RowCountChangeClass(
                    type=change_type,
                    operator=definition.operator.value,
                    parameters=VolumeAssertionDefinition._build_assertion_parameters(
                        definition.operator, definition.parameters
                    ),
                ),
            )
            if filter is not None:
                volume_info.filter = filter
            return volume_info
        else:
            raise SDKUsageError(
                f"Unsupported volume assertion definition type: {type(definition)}"
            )

    @staticmethod
    def _format_number_value(value: Union[int, float]) -> str:
        """Format number value for DataHub parameter strings.

        Converts whole numbers to integers (100.0 -> "100") and keeps decimals (100.5 -> "100.5").
        """
        if isinstance(value, float) and value.is_integer():
            return str(int(value))
        return str(value)

    @staticmethod
    def _build_assertion_parameters(
        operator: VolumeAssertionOperator,
        parameters: Union[float, Tuple[float, float]],
    ) -> models.AssertionStdParametersClass:
        """Build assertion parameters for DataHub model classes.

        Args:
            operator: The volume assertion operator.
            parameters: The parameters (int for single value, tuple for BETWEEN).

        Returns:
            AssertionStdParametersClass with appropriate parameter structure.
        """
        if operator == VolumeAssertionOperator.BETWEEN:
            assert isinstance(parameters, tuple) and len(parameters) == 2, (
                f"BETWEEN operator requires tuple of two numbers, got: {parameters}"
            )
            # Sort values to ensure minValue is actually the minimum and maxValue is the maximum
            min_val, max_val = sorted(parameters)
            return models.AssertionStdParametersClass(
                minValue=models.AssertionStdParameterClass(
                    value=VolumeAssertionDefinition._format_number_value(min_val),
                    type=models.AssertionStdParameterTypeClass.NUMBER,
                ),
                maxValue=models.AssertionStdParameterClass(
                    value=VolumeAssertionDefinition._format_number_value(max_val),
                    type=models.AssertionStdParameterTypeClass.NUMBER,
                ),
            )
        else:
            # Single value operators
            assert isinstance(parameters, (int, float)), (
                f"Single value operator {operator} requires number parameter, got: {parameters}"
            )
            return models.AssertionStdParametersClass(
                value=models.AssertionStdParameterClass(
                    value=VolumeAssertionDefinition._format_number_value(parameters),
                    type=models.AssertionStdParameterTypeClass.NUMBER,
                ),
            )

    @staticmethod
    def _extract_volume_parameters(
        assertion_urn: str,
        operator: VolumeAssertionOperator,
        parameters: models.AssertionStdParametersClass,
    ) -> Union[float, Tuple[float, float]]:
        """Extract parameters from assertion based on operator type."""
        if operator.value == "BETWEEN":
            if parameters.minValue is None or parameters.maxValue is None:
                raise SDKNotYetSupportedError(
                    f"Volume assertion {assertion_urn} has BETWEEN operator but missing min/max values"
                )
            return (float(parameters.minValue.value), float(parameters.maxValue.value))
        else:
            if parameters.value is None:
                raise SDKNotYetSupportedError(
                    f"Volume assertion {assertion_urn} has {operator.value} operator but missing value"
                )
            return float(parameters.value.value)

    @staticmethod
    def _get_row_count_total(assertion: Assertion) -> RowCountTotal:
        """Extract RowCountTotal from assertion."""
        assert isinstance(assertion.info, models.VolumeAssertionInfoClass)
        if assertion.info.rowCountTotal is None:
            raise SDKNotYetSupportedError(
                f"Volume assertion {assertion.urn} has ROW_COUNT_TOTAL type but no rowCountTotal, which is not supported"
            )
        row_count_total = assertion.info.rowCountTotal
        operator = VolumeAssertionOperator(row_count_total.operator)
        parameters = VolumeAssertionDefinition._extract_volume_parameters(
            str(assertion.urn), operator, row_count_total.parameters
        )
        return RowCountTotal(operator=operator, parameters=parameters)

    @staticmethod
    def _get_row_count_change(assertion: Assertion) -> RowCountChange:
        """Extract RowCountChange from assertion."""
        assert isinstance(assertion.info, models.VolumeAssertionInfoClass)
        if assertion.info.rowCountChange is None:
            raise SDKNotYetSupportedError(
                f"Volume assertion {assertion.urn} has ROW_COUNT_CHANGE type but no rowCountChange, which is not supported"
            )
        row_count_change = assertion.info.rowCountChange
        operator = VolumeAssertionOperator(row_count_change.operator)
        parameters = VolumeAssertionDefinition._extract_volume_parameters(
            str(assertion.urn), operator, row_count_change.parameters
        )
        kind: Literal["absolute", "percent"] = (
            "absolute"
            if row_count_change.type == models.AssertionValueChangeTypeClass.ABSOLUTE
            else "percent"
        )
        return RowCountChange(operator=operator, parameters=parameters, kind=kind)

    @staticmethod
    def from_assertion(assertion: Assertion) -> _VolumeAssertionDefinitionTypes:
        """Create a volume assertion definition from a DataHub assertion entity.

        Args:
            assertion: The DataHub assertion entity to extract the definition from.

        Returns:
            A volume assertion definition object (RowCountTotal or RowCountChange).

        Raises:
            SDKNotYetSupportedError: If the assertion is not a volume assertion or has
                unsupported configuration.
        """
        if assertion.info is None:
            raise SDKNotYetSupportedError(
                f"Assertion {assertion.urn} does not have a volume assertion info, which is not supported"
            )
        if not isinstance(assertion.info, models.VolumeAssertionInfoClass):
            raise SDKNotYetSupportedError(
                f"Assertion {assertion.urn} is not a volume assertion"
            )

        if assertion.info.type == models.VolumeAssertionTypeClass.ROW_COUNT_TOTAL:
            return VolumeAssertionDefinition._get_row_count_total(assertion)
        elif assertion.info.type == models.VolumeAssertionTypeClass.ROW_COUNT_CHANGE:
            return VolumeAssertionDefinition._get_row_count_change(assertion)
        else:
            raise SDKNotYetSupportedError(
                f"Volume assertion {assertion.urn} has unsupported type {assertion.info.type}"
            )


class _VolumeAssertionInput(_AssertionInput):
    def __init__(
        self,
        *,
        # Required fields
        dataset_urn: Union[str, DatasetUrn],
        entity_client: EntityClient,  # Needed to get the schema field spec for the detection mechanism if needed
        urn: Optional[Union[str, AssertionUrn]] = None,
        # Optional fields
        display_name: Optional[str] = None,
        enabled: bool = True,
        schedule: Optional[Union[str, models.CronScheduleClass]] = None,
        detection_mechanism: DetectionMechanismInputTypes = None,
        incident_behavior: Optional[
            Union[AssertionIncidentBehavior, list[AssertionIncidentBehavior]]
        ] = None,
        tags: Optional[TagsInputType] = None,
        created_by: Union[str, CorpUserUrn],
        created_at: datetime,
        updated_by: Union[str, CorpUserUrn],
        updated_at: datetime,
        # volume assertion fields
        definition: Optional[
            VolumeAssertionDefinitionInputTypes
        ] = None,  # TBC: default value does not make sense
    ):
        _AssertionInput.__init__(
            self,
            dataset_urn=dataset_urn,
            entity_client=entity_client,
            urn=urn,
            display_name=display_name,
            enabled=enabled,
            schedule=schedule,
            detection_mechanism=detection_mechanism,
            incident_behavior=incident_behavior,
            tags=tags,
            source_type=models.AssertionSourceTypeClass.NATIVE,  # Native assertions are of type native, not inferred
            created_by=created_by,
            created_at=created_at,
            updated_by=updated_by,
            updated_at=updated_at,
        )

        if definition is None:
            raise SDKUsageError("Volume assertion definition is required")
        self.definition = VolumeAssertionDefinition.parse(definition)

    def _assertion_type(self) -> str:
        return models.AssertionTypeClass.VOLUME

    def _create_assertion_info(
        self, filter: Optional[models.DatasetFilterClass]
    ) -> AssertionInfoInputType:
        """
        Create a VolumeAssertionInfoClass for a volume assertion.

        Args:
            filter: Optional filter to apply to the assertion.

        Returns:
            A VolumeAssertionInfoClass configured for volume assertions.
        """
        return VolumeAssertionDefinition.build_model_volume_info(
            self.definition, str(self.dataset_urn), filter
        )

    def _create_monitor_info(
        self,
        assertion_urn: AssertionUrn,
        status: models.MonitorStatusClass,
        schedule: models.CronScheduleClass,
    ) -> models.MonitorInfoClass:
        """
        Create a MonitorInfoClass with all the necessary components.
        """
        source_type, field = self._convert_assertion_source_type_and_field()
        return models.MonitorInfoClass(
            type=models.MonitorTypeClass.ASSERTION,
            status=status,
            assertionMonitor=models.AssertionMonitorClass(
                assertions=[
                    models.AssertionEvaluationSpecClass(
                        assertion=str(assertion_urn),
                        schedule=schedule,
                        parameters=self._get_assertion_evaluation_parameters(
                            str(source_type), field
                        ),
                    )
                ]
            ),
        )

    def _convert_schedule(self) -> models.CronScheduleClass:
        """Create a schedule for a volume assertion.

        Returns:
            A CronScheduleClass with appropriate schedule settings.
        """
        if self.schedule is None:
            return DEFAULT_EVERY_SIX_HOURS_SCHEDULE

        return models.CronScheduleClass(
            cron=self.schedule.cron,
            timezone=self.schedule.timezone,
        )

    def _get_assertion_evaluation_parameters(
        self, source_type: str, field: Optional[FieldSpecType]
    ) -> models.AssertionEvaluationParametersClass:
        return models.AssertionEvaluationParametersClass(
            type=models.AssertionEvaluationParametersTypeClass.DATASET_VOLUME,
            datasetVolumeParameters=models.DatasetVolumeAssertionParametersClass(
                sourceType=source_type
            ),
        )

    def _convert_assertion_source_type_and_field(
        self,
    ) -> tuple[str, Optional[FieldSpecType]]:
        """Convert the detection mechanism to source type and field."""
        default_source_type = models.DatasetVolumeSourceTypeClass.INFORMATION_SCHEMA

        if self.detection_mechanism is None:
            return default_source_type, None

        # Convert detection mechanism to volume source type
        if isinstance(self.detection_mechanism, str):
            if self.detection_mechanism == "information_schema":
                return models.DatasetVolumeSourceTypeClass.INFORMATION_SCHEMA, None
            elif self.detection_mechanism == "datahub_operation":
                return models.DatasetVolumeSourceTypeClass.OPERATION, None
            else:
                return default_source_type, None

        # For more complex detection mechanisms, we might need additional logic
        return default_source_type, None
