from __future__ import annotations

import copyreg
import dataclasses
import datetime
import decimal
import hashlib
import inspect
import pathlib
import pickle
import threading
import types
import uuid
from collections import OrderedDict
from collections.abc import Iterable
from contextvars import ContextVar
from functools import partial

import cloudpickle
from tlz import curry, identity
from tlz.functoolz import Compose

from dask import config
from dask.core import literal
from dask.hashing import hash_buffer_hex
from dask.utils import Dispatch


class TokenizationError(RuntimeError):
    pass


def _tokenize(*args: object, **kwargs: object) -> str:
    token: object = _normalize_seq_func(args)
    if kwargs:
        token = token, _normalize_seq_func(sorted(kwargs.items()))

    # Pass `usedforsecurity=False` to support FIPS builds of Python
    return hashlib.md5(str(token).encode(), usedforsecurity=False).hexdigest()


tokenize_lock = threading.RLock()
_SEEN: dict[int, tuple[int, object]] = {}
_ENSURE_DETERMINISTIC: ContextVar[bool | None] = ContextVar("_ENSURE_DETERMINISTIC")


def tokenize(
    *args: object, ensure_deterministic: bool | None = None, **kwargs: object
) -> str:
    """Deterministic token

    >>> tokenize([1, 2, '3'])  # doctest: +SKIP
    '06961e8de572e73c2e74b51348177918'

    >>> tokenize('Hello') == tokenize('Hello')
    True

    Parameters
    ----------
    args, kwargs:
        objects to tokenize
    ensure_deterministic: bool, optional
        If True, raise TokenizationError if the objects cannot be deterministically
        tokenized, e.g. two identical objects will return different tokens.
        Defaults to the `tokenize.ensure-deterministic` configuration parameter.
    """
    global _SEEN
    with tokenize_lock:
        seen_before, _SEEN = _SEEN, {}
        token = None
        try:
            _ENSURE_DETERMINISTIC.get()
        except LookupError:
            token = _ENSURE_DETERMINISTIC.set(ensure_deterministic)
        try:
            return _tokenize(*args, **kwargs)
        finally:
            if token:
                _ENSURE_DETERMINISTIC.reset(token)
            _SEEN = seen_before


def _maybe_raise_nondeterministic(msg: str) -> None:
    try:
        val = _ENSURE_DETERMINISTIC.get()
    except LookupError:
        val = None
    if val or val is None and config.get("tokenize.ensure-deterministic"):
        raise TokenizationError(msg)


_IDENTITY_DISPATCH = (
    int,
    float,
    str,
    bytes,
    type(None),
    slice,
    complex,
    type(Ellipsis),
    decimal.Decimal,
    datetime.date,
    datetime.time,
    datetime.datetime,
    datetime.timedelta,
    pathlib.PurePath,
)
normalize_token = Dispatch()
normalize_token.register(
    _IDENTITY_DISPATCH,
    identity,
)


@normalize_token.register((types.MappingProxyType, dict))
def normalize_dict(d):
    with tokenize_lock:
        if id(d) in _SEEN:
            return "__seen", _SEEN[id(d)][0]
        _SEEN[id(d)] = len(_SEEN), d
        try:
            return "dict", _normalize_seq_func(
                sorted(d.items(), key=lambda kv: str(kv[0]))
            )
        finally:
            _SEEN.pop(id(d), None)


@normalize_token.register(OrderedDict)
def normalize_ordered_dict(d):
    return _normalize_seq_func((type(d), list(d.items())))


@normalize_token.register(set)
def normalize_set(s):
    # Note: in some Python version / OS combinations, set order changes every
    # time you recreate the set (even within the same interpreter).
    # In most other cases, set ordering is consistent within the same interpreter.
    return "set", _normalize_seq_func(sorted(s, key=str))


def _normalize_seq_func(seq: Iterable[object]) -> tuple[object, ...]:
    def _inner_normalize_token(item):
        # Don't go through Dispatch. That's slow
        if isinstance(item, _IDENTITY_DISPATCH):
            return item
        return normalize_token(item)

    with tokenize_lock:
        if id(seq) in _SEEN:
            return "__seen", _SEEN[id(seq)][0]
        _SEEN[id(seq)] = len(_SEEN), seq
        try:
            return tuple(map(_inner_normalize_token, seq))
        finally:
            del _SEEN[id(seq)]


@normalize_token.register((tuple, list))
def normalize_seq(seq):
    return type(seq).__name__, _normalize_seq_func(seq)


@normalize_token.register(literal)
def normalize_literal(lit):
    return "literal", normalize_token(lit())


@normalize_token.register(Compose)
def normalize_compose(func):
    return _normalize_seq_func((func.first,) + func.funcs)


@normalize_token.register((partial, curry))
def normalize_partial(func):
    return _normalize_seq_func((func.func, func.args, func.keywords))


@normalize_token.register((types.MethodType, types.MethodWrapperType))
def normalize_bound_method(meth):
    return normalize_token(meth.__self__), meth.__name__


@normalize_token.register(types.BuiltinFunctionType)
def normalize_builtin_function_or_method(func):
    # Note: BuiltinMethodType is BuiltinFunctionType
    self = getattr(func, "__self__", None)
    if self is not None and not inspect.ismodule(self):
        return normalize_bound_method(func)
    else:
        return normalize_object(func)


@normalize_token.register(object)
def normalize_object(o):
    method = getattr(o, "__dask_tokenize__", None)
    if method is not None and not isinstance(o, type):
        return method()

    if type(o) is object:
        return _normalize_pure_object(o)

    if isinstance(o, type):
        copyreg._slotnames(o)

    if dataclasses.is_dataclass(o) and not isinstance(o, type):
        return _normalize_dataclass(o)

    try:
        return _normalize_pickle(o)
    except Exception:
        _maybe_raise_nondeterministic(
            f"Object {o!r} cannot be deterministically hashed. This likely "
            "indicates that the object cannot be serialized deterministically."
        )
        return uuid.uuid4().hex


_seen_objects = set()


def _normalize_pure_object(o: object) -> tuple[str, int]:
    _maybe_raise_nondeterministic(
        "object() cannot be deterministically hashed. See "
        "https://docs.dask.org/en/latest/custom-collections.html#implementing-deterministic-hashing "
        "for more information."
    )
    # Idempotent, but not deterministic. Make sure that the id is not reused.
    _seen_objects.add(o)
    return "object", id(o)


def _normalize_pickle(o: object) -> tuple:
    buffers: list[pickle.PickleBuffer] = []
    pik: int | None = None
    pik2: int | None = None
    for _ in range(3):
        buffers.clear()
        try:
            out = pickle.dumps(o, protocol=5, buffer_callback=buffers.append)
            if b"__main__" in out:
                # Use `cloudpickle` for objects defined in `__main__`
                buffers.clear()
                out = cloudpickle.dumps(o, protocol=5, buffer_callback=buffers.append)
            pickle.loads(out, buffers=buffers)
            pik2 = hash_buffer_hex(out)
        except Exception:
            buffers.clear()
            try:
                out = cloudpickle.dumps(o, protocol=5, buffer_callback=buffers.append)
                pickle.loads(out, buffers=buffers)
                pik2 = hash_buffer_hex(out)
            except Exception:
                break
        if pik and pik2 and pik == pik2:
            break
        pik = pik2
    else:
        _maybe_raise_nondeterministic("Failed to tokenize deterministically")
    if pik is None:
        _maybe_raise_nondeterministic("Failed to tokenize deterministically")
        pik = int(uuid.uuid4())
    return pik, [hash_buffer_hex(buf) for buf in buffers]


def _normalize_dataclass(obj):
    fields = [
        (field.name, normalize_token(getattr(obj, field.name, None)))
        for field in dataclasses.fields(obj)
    ]
    params = obj.__dataclass_params__
    params = [(attr, getattr(params, attr)) for attr in params.__slots__]

    return normalize_object(type(obj)), params, fields


@normalize_token.register_lazy("pandas")
def register_pandas():
    import pandas as pd

    # use dask._pandas_compat to avoid importing dask.dataframe here
    from dask._pandas_compat import PANDAS_GE_210

    @normalize_token.register(pd.RangeIndex)
    def normalize_range_index(x):
        return type(x), x.start, x.stop, x.step, x.dtype, x.name

    @normalize_token.register(pd.Index)
    def normalize_index(ind):
        values = ind.array

        if isinstance(values, pd.arrays.ArrowExtensionArray):
            import pyarrow as pa

            # these are sensitive to fragmentation of the backing Arrow array.
            # Because common operations like DataFrame.getitem and DataFrame.setitem
            # result in fragmented Arrow arrays, we'll consolidate them here.

            if PANDAS_GE_210:
                # avoid combining chunks by using chunked_array
                values = pa.chunked_array([values._pa_array]).combine_chunks()
            else:
                values = pa.array(values)

        return type(ind), ind.name, normalize_token(values)

    @normalize_token.register(pd.MultiIndex)
    def normalize_index(ind):
        codes = ind.codes
        return (
            [ind.name]
            + [normalize_token(x) for x in ind.levels]
            + [normalize_token(x) for x in codes]
        )

    @normalize_token.register(pd.Categorical)
    def normalize_categorical(cat):
        return [normalize_token(cat.codes), normalize_token(cat.dtype)]

    @normalize_token.register(pd.arrays.PeriodArray)
    @normalize_token.register(pd.arrays.DatetimeArray)
    @normalize_token.register(pd.arrays.TimedeltaArray)
    def normalize_period_array(arr):
        return [normalize_token(arr.asi8), normalize_token(arr.dtype)]

    @normalize_token.register(pd.arrays.IntervalArray)
    def normalize_interval_array(arr):
        return [
            normalize_token(arr.left),
            normalize_token(arr.right),
            normalize_token(arr.closed),
        ]

    @normalize_token.register(pd.Series)
    def normalize_series(s):
        return [
            s.name,
            s.dtype,
            normalize_token(s._values),
            normalize_token(s.index),
        ]

    @normalize_token.register(pd.DataFrame)
    def normalize_dataframe(df):
        mgr = df._mgr
        data = list(mgr.arrays) + [df.columns, df.index]
        return list(map(normalize_token, data))

    @normalize_token.register(pd.arrays.ArrowExtensionArray)
    def normalize_extension_array(arr):
        try:
            return (type(arr), normalize_token(arr._pa_array))
        except AttributeError:
            return (type(arr), normalize_token(arr._data))

    @normalize_token.register(pd.api.extensions.ExtensionArray)
    def normalize_extension_array(arr):
        import numpy as np

        return normalize_token(np.asarray(arr))

    # Dtypes
    @normalize_token.register(pd.api.types.CategoricalDtype)
    def normalize_categorical_dtype(dtype):
        return [normalize_token(dtype.categories), normalize_token(dtype.ordered)]

    @normalize_token.register(pd.api.extensions.ExtensionDtype)
    def normalize_period_dtype(dtype):
        return normalize_token(dtype.name)

    @normalize_token.register(type(pd.NA))
    def normalize_na(na):
        return pd.NA

    @normalize_token.register(pd.offsets.BaseOffset)
    def normalize_offset(offset):
        return offset.freqstr


@normalize_token.register_lazy("numba")
def register_numba():
    import numba

    @normalize_token.register(numba.core.serialize.ReduceMixin)
    def normalize_numba_ufunc(obj):
        return normalize_token((obj._reduce_class(), obj._reduce_states()))


@normalize_token.register_lazy("pyarrow")
def register_pyarrow():
    import pyarrow as pa

    @normalize_token.register(pa.DataType)
    def normalize_datatype(dt):
        return pickle.dumps(dt, protocol=4)

    @normalize_token.register(pa.Table)
    def normalize_table(dt):
        return (
            "pa.Table",
            normalize_token(dt.schema),
            normalize_token(dt.columns),
        )

    @normalize_token.register(pa.ChunkedArray)
    def normalize_chunked_array(arr):
        return (
            "pa.ChunkedArray",
            normalize_token(arr.type),
            normalize_token(arr.chunks),
        )

    @normalize_token.register(pa.Array)
    def normalize_array(arr):
        buffers = arr.buffers()
        # pyarrow does something clever when (de)serializing an array that has
        # an empty validity map: The buffers for the deserialized array will
        # have `None` instead of the empty validity map.
        #
        # We'll replicate that behavior here to ensure we get consistent
        # tokenization.
        buffers = arr.buffers()
        if len(buffers) and buffers[0] is not None and arr.null_count == 0:
            buffers[0] = None

        return (
            "pa.Array",
            normalize_token(arr.type),
            normalize_token(buffers),
        )

    @normalize_token.register(pa.Buffer)
    def normalize_buffer(buf):
        return ("pa.Buffer", hash_buffer_hex(buf))


@normalize_token.register_lazy("numpy")
def register_numpy():
    import numpy as np

    @normalize_token.register(np.ndarray)
    def normalize_array(x):
        if not x.shape:
            return (x.item(), x.dtype)
        if x.dtype.hasobject:
            try:
                try:
                    # string fast-path
                    data = hash_buffer_hex(
                        "-".join(x.flat).encode(
                            encoding="utf-8", errors="surrogatepass"
                        )
                    )
                except UnicodeDecodeError:
                    # bytes fast-path
                    data = hash_buffer_hex(b"-".join(x.flat))
            except (TypeError, UnicodeDecodeError):
                return normalize_object(x)
        else:
            try:
                data = hash_buffer_hex(x.ravel(order="K").view("i1"))
            except (BufferError, AttributeError, ValueError):
                data = hash_buffer_hex(x.copy().ravel(order="K").view("i1"))
        return (data, x.dtype, x.shape)

    @normalize_token.register(np.memmap)
    def normalize_mmap(mm):
        return hash_buffer_hex(np.ascontiguousarray(mm))

    @normalize_token.register(np.ufunc)
    def normalize_ufunc(func):
        try:
            return _normalize_pickle(func)
        except Exception:
            _maybe_raise_nondeterministic(
                f"Cannot tokenize numpy ufunc {func!r}. Please use functions "
                "of the dask.array.ufunc module instead. See also "
                "https://docs.dask.org/en/latest/array-numpy-compatibility.html"
            )
            return uuid.uuid4().hex

    @normalize_token.register(np.dtype)
    def normalize_dtype(dtype):
        return dtype.str


def _tokenize_deterministic(*args, **kwargs) -> str:
    # Utility to be strict about deterministic tokens
    return tokenize(*args, ensure_deterministic=True, **kwargs)
