"""
Classes and functions to treat column data types according to
Deirokay data types.
"""

from typing import Iterable, List, Optional, Union

import dask.dataframe  # lazy module
import numpy  # lazy module
import pandas  # lazy module

from deirokay._typing import DeirokaySerializedSeries
from deirokay.enums import Backend, DTypes

from ..multibackend import serialize, treat
from .validator import Validator


class BooleanTreater(Validator):
    """Treater for boolean-like variables

    Parameters
    ----------
    truthies : List[str], optional
        List of strings to be considered as True,
        by default ['true', 'True']
    falsies : List[str], optional
        List of strings to be considered as False,
        by default ['false', 'False']
    ignore_case : bool, optional
        Ignore case when evaluating list of values,
        by default False
    default_value : Optional[bool], optional
        Value to replace null values (None, pandas.NA or numpy,NaN).
        By default None.
    """

    supported_backends = [Backend.PANDAS, Backend.DASK]
    supported_dtype = DTypes.BOOLEAN
    supported_primitives = [bool]

    def __init__(
        self,
        truthies: List[str] = ["true", "True"],
        falsies: List[str] = ["false", "False"],
        ignore_case: bool = False,
        default_value: Optional[bool] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)

        assert default_value in (True, False, None)

        self.ignore_case = ignore_case
        self.default_value = default_value

        if self.ignore_case:
            self.truthies = set(truthy.lower() for truthy in truthies)
            self.falsies = set(falsy.lower() for falsy in falsies)
        else:
            self.truthies = set(truthies)
            self.falsies = set(falsies)

        if self.truthies & self.falsies:
            raise ValueError("Truthies and Falsies sets should be" " disjoint sets.")

    def _evaluate(self, value: Union[bool, str, None]) -> Union[bool, None]:
        if value is True:
            return True
        if value is False:
            return False
        if value is None or value is numpy.nan or value is pandas.NA:
            return self.default_value

        _value = value.lower() if self.ignore_case else value
        if _value in self.truthies:
            return True
        if _value in self.falsies:
            return False
        raise ValueError(
            f'Unexpected boolean value: "{value}"'
            f" ({value.__class__})\n"
            f"Expected values: {self.truthies | self.falsies}"
        )

    @treat(Backend.PANDAS)
    def _treat_pandas(self, series: Iterable) -> "pandas.Series":
        series = super()._treat_pandas(series)
        series = series.apply(self._evaluate).astype("boolean")
        # Validate again
        super()._treat_pandas(series)

        return series

    @treat(Backend.DASK)
    def _treat_dask(self, series: Iterable) -> "dask.dataframe.Series":
        series = super()._treat_dask(series)
        series = series.apply(self._evaluate, meta=(series.name, "object")).astype(
            "boolean"
        )
        # Validate again
        super()._treat_dask(series)

        return series

    @staticmethod
    def _serialize_common(series):
        def _convert(item):
            if item is pandas.NA:
                return None
            return bool(item)

        return {
            "values": [_convert(item) for item in series],
            "parser": {"dtype": BooleanTreater.supported_dtype.value},
        }

    @serialize(Backend.PANDAS)
    @staticmethod
    def _serialize_pandas(series: "pandas.Series") -> DeirokaySerializedSeries:
        return BooleanTreater._serialize_common(series)

    @serialize(Backend.DASK)
    @staticmethod
    def _serialize_dask(series: "dask.dataframe.Series") -> DeirokaySerializedSeries:
        return BooleanTreater._serialize_common(series)
