# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2026)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Hashing for st.cache_data and st.cache_resource."""

from __future__ import annotations

import collections
import collections.abc
import dataclasses
import datetime
import functools
import inspect
import io
import os
import pickle  # noqa: S403
import re
import sys
import tempfile
import threading
import uuid
import weakref
from collections.abc import Callable
from enum import Enum
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, TypeAlias, cast

from streamlit import logger, type_util, util
from streamlit.errors import StreamlitAPIException
from streamlit.runtime.caching.cache_errors import UnhashableTypeError
from streamlit.runtime.caching.cache_type import CacheType
from streamlit.runtime.uploaded_file_manager import UploadedFile

if TYPE_CHECKING:
    import numpy.typing as npt
    import pandas as pd
    from PIL.Image import Image

_LOGGER: Final = logger.get_logger(__name__)

# If a dataframe has more than this many rows, we consider it large and hash a sample.
_PANDAS_ROWS_LARGE: Final = 50_000
_PANDAS_SAMPLE_SIZE: Final = 10_000

# Similar to dataframes, we also sample large numpy arrays.
_NP_SIZE_LARGE: Final = 500_000
_NP_SAMPLE_SIZE: Final = 100_000

HashFuncsDict: TypeAlias = dict[str | type[Any], Callable[[Any], Any]]

# Arbitrary item to denote where we found a cycle in a hashed object.
# This allows us to hash self-referencing lists, dictionaries, etc.
_CYCLE_PLACEHOLDER: Final = (
    b"streamlit-57R34ML17-hesamagicalponyflyingthroughthesky-CYCLE"
)


class UserHashError(StreamlitAPIException):
    def __init__(
        self,
        orig_exc: BaseException,
        object_to_hash: Any,
        hash_func: Callable[[Any], Any],
        cache_type: CacheType | None = None,
    ) -> None:
        self.alternate_name = type(orig_exc).__name__
        self.hash_func = hash_func
        self.cache_type = cache_type

        msg = self._get_message_from_func(orig_exc, object_to_hash)

        super().__init__(msg)
        self.with_traceback(orig_exc.__traceback__)

    def _get_message_from_func(
        self,
        orig_exc: BaseException,
        cached_func: Any,
    ) -> str:
        args = self._get_error_message_args(orig_exc, cached_func)

        return (
            f"""
{args["orig_exception_desc"]}

This error is likely due to a bug in {args["hash_func_name"]}, which is a
user-defined hash function that was passed into the `{args["cache_primitive"]}` decorator of
{args["object_desc"]}.

{args["hash_func_name"]} failed when hashing an object of type
`{args["failed_obj_type_str"]}`.  If you don't know where that object is coming from,
try looking at the hash chain below for an object that you do recognize, then
pass that to `hash_funcs` instead:

```
{args["hash_stack"]}
```

If you think this is actually a Streamlit bug, please
[file a bug report here](https://github.com/streamlit/streamlit/issues/new/choose).
"""
        ).strip("\n")

    def _get_error_message_args(
        self,
        orig_exc: BaseException,
        failed_obj: Any,
    ) -> dict[str, Any]:
        hash_source = hash_stacks.current.hash_source

        failed_obj_type_str = type_util.get_fqn_type(failed_obj)

        if hash_source is None:
            object_desc = "something"
        elif hasattr(hash_source, "__name__"):
            object_desc = f"`{hash_source.__name__}()`"
        else:
            object_desc = "a function"

        decorator_name = ""
        if self.cache_type is CacheType.RESOURCE:
            decorator_name = "@st.cache_resource"
        elif self.cache_type is CacheType.DATA:
            decorator_name = "@st.cache_data"

        hash_func_name = (
            f"`{self.hash_func.__name__}()`"
            if hasattr(self.hash_func, "__name__")
            else "a function"
        )

        return {
            "orig_exception_desc": str(orig_exc),
            "failed_obj_type_str": failed_obj_type_str,
            "hash_stack": hash_stacks.current.pretty_print(),
            "object_desc": object_desc,
            "cache_primitive": decorator_name,
            "hash_func_name": hash_func_name,
        }


def update_hash(
    val: Any,
    hasher: Any,
    cache_type: CacheType,
    hash_source: Callable[..., Any] | None = None,
    hash_funcs: HashFuncsDict | None = None,
) -> None:
    """Updates a hashlib hasher with the hash of val.

    This is the main entrypoint to hashing.py.
    """

    hash_stacks.current.hash_source = hash_source

    ch = _CacheFuncHasher(cache_type, hash_funcs)
    ch.update(hasher, val)


class _HashStack:
    """Stack of what has been hashed, for debug and circular reference detection.

    This internally keeps 1 stack per thread.

    Internally, this stores the ID of pushed objects rather than the objects
    themselves because otherwise the "in" operator inside __contains__ would
    fail for objects that don't return a boolean for "==" operator. For
    example, arr == 10 where arr is a NumPy array returns another NumPy array.
    This causes the "in" to crash since it expects a boolean.
    """

    def __init__(self) -> None:
        self._stack: collections.OrderedDict[int, list[Any]] = collections.OrderedDict()
        # A function that we decorate with streamlit cache
        # primitive (st.cache_data or st.cache_resource).
        self.hash_source: Callable[..., Any] | None = None

    def __repr__(self) -> str:
        return util.repr_(self)

    def push(self, val: Any) -> None:
        self._stack[id(val)] = val

    def pop(self) -> None:
        self._stack.popitem()

    def __contains__(self, val: Any) -> bool:
        return id(val) in self._stack

    def pretty_print(self) -> str:
        def to_str(v: Any) -> str:
            try:
                return f"Object of type {type_util.get_fqn_type(v)}: {v}"
            except Exception:
                return "<Unable to convert item to string>"

        return "\n".join(to_str(x) for x in reversed(self._stack.values()))


class _HashStacks:
    """Stacks of what has been hashed, with at most 1 stack per thread."""

    def __init__(self) -> None:
        self._stacks: weakref.WeakKeyDictionary[threading.Thread, _HashStack] = (
            weakref.WeakKeyDictionary()
        )

    def __repr__(self) -> str:
        return util.repr_(self)

    @property
    def current(self) -> _HashStack:
        current_thread = threading.current_thread()

        stack = self._stacks.get(current_thread, None)

        if stack is None:
            stack = _HashStack()
            self._stacks[current_thread] = stack

        return stack


hash_stacks = _HashStacks()


def _int_to_bytes(i: int) -> bytes:
    num_bytes = (i.bit_length() + 8) // 8
    return i.to_bytes(num_bytes, "little", signed=True)


def _float_to_bytes(f: float) -> bytes:
    # Lazy-load for performance reasons.
    import struct

    # Floats are 64bit in Python, so we need to use the "d" format.
    return struct.pack("<d", f)


def _key(obj: Any | None) -> Any:
    """Return key for memoization."""

    if obj is None:
        return None

    def is_simple(obj: Any) -> bool:
        return (
            isinstance(obj, (bytes, bytearray, str, float, int, bool, uuid.UUID))
            or obj is None
        )

    if is_simple(obj):
        return obj

    if isinstance(obj, tuple) and all(map(is_simple, obj)):
        return obj

    if isinstance(obj, list) and all(map(is_simple, obj)):
        return ("__l", tuple(obj))

    if inspect.isbuiltin(obj) or inspect.isroutine(obj) or inspect.iscode(obj):
        return id(obj)

    return NoResult


class _CacheFuncHasher:
    """A hasher that can hash objects with cycles."""

    def __init__(
        self, cache_type: CacheType, hash_funcs: HashFuncsDict | None = None
    ) -> None:
        # Can't use types as the keys in the internal _hash_funcs because
        # we always remove user-written modules from memory when rerunning a
        # script in order to reload it and grab the latest code changes.
        # (See LocalSourcesWatcher.py:on_file_changed) This causes
        # the type object to refer to different underlying class instances each run,
        # so type-based comparisons fail. To solve this, we use the types converted
        # to fully-qualified strings as keys in our internal dict.
        self._hash_funcs: HashFuncsDict
        if hash_funcs:
            self._hash_funcs = {
                k if isinstance(k, str) else type_util.get_fqn(k): v
                for k, v in hash_funcs.items()
            }
        else:
            self._hash_funcs = {}
        self._hashes: dict[Any, bytes] = {}

        # The number of the bytes in the hash.
        self.size = 0

        self.cache_type = cache_type

    def __repr__(self) -> str:
        return util.repr_(self)

    def to_bytes(self, obj: Any) -> bytes:
        """Add memoization to _to_bytes and protect against cycles in data structures."""
        tname = type(obj).__qualname__.encode()
        key = (tname, _key(obj))

        # Memoize if possible.
        if key[1] is not NoResult and key in self._hashes:
            return self._hashes[key]

        # Break recursive cycles.
        if obj in hash_stacks.current:
            return _CYCLE_PLACEHOLDER

        hash_stacks.current.push(obj)

        try:
            # Hash the input
            b = b"%s:%s" % (tname, self._to_bytes(obj))

            # Hmmm... It's possible that the size calculation is wrong. When we
            # call to_bytes inside _to_bytes things get double-counted.
            self.size += sys.getsizeof(b)

            if key[1] is not NoResult:
                self._hashes[key] = b

        finally:
            # In case an UnhashableTypeError (or other) error is thrown, clean up the
            # stack so we don't get false positives in future hashing calls
            hash_stacks.current.pop()

        return b

    def update(self, hasher: Any, obj: Any) -> None:
        """Update the provided hasher with the hash of an object."""
        b = self.to_bytes(obj)
        hasher.update(b)

    def _to_bytes(self, obj: Any) -> bytes:
        """Hash objects to bytes, including code with dependencies.

        Python's built in `hash` does not produce consistent results across
        runs.
        """

        h = util.create_fast_hasher()

        if type_util.is_type(obj, "unittest.mock.Mock") or type_util.is_type(
            obj, "unittest.mock.MagicMock"
        ):
            # Mock objects can appear to be infinitely
            # deep, so we don't try to hash them at all.
            return self.to_bytes(id(obj))

        if isinstance(obj, (bytes, bytearray)):
            return obj

        if type_util.get_fqn_type(obj) in self._hash_funcs:
            # Escape hatch for unsupported objects
            hash_func = self._hash_funcs[type_util.get_fqn_type(obj)]
            try:
                output = hash_func(obj)
            except Exception as ex:
                raise UserHashError(
                    ex, obj, hash_func=hash_func, cache_type=self.cache_type
                ) from ex
            return self.to_bytes(output)

        if isinstance(obj, str):
            return obj.encode()

        if isinstance(obj, float):
            return _float_to_bytes(obj)

        if isinstance(obj, int):
            return _int_to_bytes(obj)

        if isinstance(obj, uuid.UUID):
            return obj.bytes

        if isinstance(obj, datetime.datetime):
            return obj.isoformat().encode()

        if isinstance(obj, (list, tuple)):
            for item in obj:
                self.update(h, item)
            return h.digest()

        if isinstance(obj, dict):
            for item in obj.items():
                self.update(h, item)
            return h.digest()

        if obj is None:
            return b"0"

        if obj is True:  # pragma: no cover - unreachable; bool subclasses int
            return b"1"

        if obj is False:  # pragma: no cover - unreachable; bool subclasses int
            return b"0"

        if not isinstance(obj, type) and dataclasses.is_dataclass(obj):
            return self.to_bytes(dataclasses.asdict(obj))

        if isinstance(obj, Enum):
            return str(obj).encode()

        # pandas 3.x changed __module__ from pandas.core.* to pandas.*
        if type_util.is_type(obj, re.compile(r"^pandas(\.core\.series)?\.Series$")):
            from pandas.util import hash_pandas_object

            series_obj = cast("pd.Series[Any]", obj)
            self.update(h, series_obj.size)
            self.update(h, series_obj.dtype.name)

            if len(series_obj) >= _PANDAS_ROWS_LARGE:
                series_obj = series_obj.sample(n=_PANDAS_SAMPLE_SIZE, random_state=0)

            try:
                self.update(h, hash_pandas_object(series_obj).to_numpy().tobytes())
                return h.digest()
            except TypeError:
                _LOGGER.warning(
                    "Pandas Series hash failed. Falling back to pickling the object.",
                    exc_info=True,
                )

                # Use pickle if pandas cannot hash the object for example if
                # it contains unhashable objects.
                return b"%s" % pickle.dumps(series_obj, pickle.HIGHEST_PROTOCOL)

        elif type_util.is_type(obj, re.compile(r"^pandas(\.core\.frame)?\.DataFrame$")):
            from pandas.util import hash_pandas_object

            df_obj: pd.DataFrame = cast("pd.DataFrame", obj)
            self.update(h, df_obj.shape)

            if len(df_obj) >= _PANDAS_ROWS_LARGE:
                df_obj = df_obj.sample(n=_PANDAS_SAMPLE_SIZE, random_state=0)

            try:
                column_hash_bytes = self.to_bytes(hash_pandas_object(df_obj.dtypes))
                self.update(h, column_hash_bytes)
                values_hash_bytes = self.to_bytes(hash_pandas_object(df_obj))
                self.update(h, values_hash_bytes)
                return h.digest()

            except TypeError:
                _LOGGER.warning(
                    "Pandas DataFrame hash failed. Falling back to pickling the object.",
                    exc_info=True,
                )

                # Use pickle if pandas cannot hash the object for example if
                # it contains unhashable objects.
                return b"%s" % pickle.dumps(df_obj, pickle.HIGHEST_PROTOCOL)

        elif type_util.is_type(obj, "polars.series.series.Series"):
            import polars as pl

            obj = cast("pl.Series", obj)
            self.update(h, str(obj.dtype).encode())
            self.update(h, obj.shape)

            if len(obj) >= _PANDAS_ROWS_LARGE:
                obj = obj.sample(n=_PANDAS_SAMPLE_SIZE, seed=0)

            try:
                self.update(h, obj.hash(seed=0).to_arrow().to_string().encode())
                return h.digest()

            except TypeError:
                _LOGGER.warning(
                    "Polars Series hash failed. Falling back to pickling the object.",
                    exc_info=True,
                )

                # Use pickle if polars cannot hash the object for example if
                # it contains unhashable objects.
                return b"%s" % pickle.dumps(obj, pickle.HIGHEST_PROTOCOL)

        elif type_util.is_type(obj, "polars.dataframe.frame.DataFrame"):
            import polars as pl  # noqa: TC002

            obj = cast("pl.DataFrame", obj)
            self.update(h, obj.shape)

            if len(obj) >= _PANDAS_ROWS_LARGE:
                obj = obj.sample(n=_PANDAS_SAMPLE_SIZE, seed=0)
            try:
                for c, t in obj.schema.items():
                    self.update(h, c.encode())
                    self.update(h, str(t).encode())

                values_hash_bytes = (
                    obj.hash_rows(seed=0).hash(seed=0).to_arrow().to_string().encode()
                )

                self.update(h, values_hash_bytes)
                return h.digest()

            except TypeError:
                _LOGGER.warning(
                    "Polars DataFrame hash failed. Falling back to pickling the object.",
                    exc_info=True,
                )

                # Use pickle if polars cannot hash the object for example if
                # it contains unhashable objects.
                return b"%s" % pickle.dumps(obj, pickle.HIGHEST_PROTOCOL)

        elif type_util.is_type(obj, "numpy.ndarray"):
            np_obj: npt.NDArray[Any] = cast("npt.NDArray[Any]", obj)
            self.update(h, np_obj.shape)
            self.update(h, str(np_obj.dtype))

            if np_obj.size >= _NP_SIZE_LARGE:
                import numpy as np

                state = np.random.RandomState(0)
                np_obj = state.choice(np_obj.flat, size=_NP_SAMPLE_SIZE)

            self.update(h, np_obj.tobytes())
            return h.digest()

        elif type_util.is_type(obj, "PIL.Image.Image"):
            import numpy as np

            pil_obj: Image = cast("Image", obj)

            # we don't just hash the results of obj.tobytes() because we want to use
            # the sampling logic for numpy data
            np_array = np.frombuffer(pil_obj.tobytes(), dtype="uint8")
            return self.to_bytes(np_array)

        elif inspect.isbuiltin(obj):
            return bytes(obj.__name__.encode())

        elif isinstance(obj, (MappingProxyType, collections.abc.ItemsView)):
            return self.to_bytes(dict(obj))

        elif type_util.is_type(obj, "builtins.getset_descriptor"):
            return bytes(obj.__qualname__.encode())

        elif isinstance(obj, UploadedFile):
            # UploadedFile is a BytesIO (thus IOBase) but has a name.
            # It does not have a timestamp so this must come before
            # temporary files
            self.update(h, obj.name)
            self.update(h, obj.tell())
            self.update(h, obj.getvalue())
            return h.digest()

        elif hasattr(obj, "name") and (
            # Handle temporary files used during testing
            isinstance(obj, (io.IOBase, tempfile._TemporaryFileWrapper))
        ):
            # Hash files as name + last modification date + offset.
            # NB: we're using hasattr("name") to differentiate between
            # on-disk and in-memory StringIO/BytesIO file representations.
            # That means that this condition must come *before* the next
            # condition, which just checks for StringIO/BytesIO.
            # Just to appease MyPy.
            obj_name = getattr(obj, "name", "wonthappen")
            self.update(h, obj_name)
            self.update(h, os.path.getmtime(obj_name))
            self.update(h, obj.tell())
            return h.digest()

        elif isinstance(obj, re.Pattern):
            return self.to_bytes([obj.pattern, obj.flags])

        elif isinstance(obj, (io.StringIO, io.BytesIO)):
            # Hash in-memory StringIO/BytesIO by their full contents
            # and seek position.
            self.update(h, obj.tell())
            self.update(h, obj.getvalue())
            return h.digest()

        elif type_util.is_type(obj, "numpy.ufunc"):
            # For numpy.remainder, this returns remainder.
            return bytes(obj.__name__.encode())

        elif inspect.ismodule(obj):
            # TODO: Figure out how to best show this kind of warning to the
            # user. In the meantime, show nothing. This scenario is too common,
            # so the current warning is quite annoying...
            # st.warning(('Streamlit does not support hashing modules. '
            #             'We did not hash `%s`.') % obj.__name__)
            # TODO: Hash more than just the name for internal modules.
            return self.to_bytes(obj.__name__)

        elif inspect.isclass(obj):
            # TODO: Figure out how to best show this kind of warning to the
            # user. In the meantime, show nothing. This scenario is too common,
            # (e.g. in every "except" statement) so the current warning is
            # quite annoying...
            # st.warning(('Streamlit does not support hashing classes. '
            #             'We did not hash `%s`.') % obj.__name__)
            # TODO: Hash more than just the name of classes.
            return self.to_bytes(obj.__name__)

        elif isinstance(obj, functools.partial):
            # The return value of functools.partial is not a plain function:
            # it's a callable object that remembers the original function plus
            # the values you pickled into it. So here we need to special-case it.
            self.update(h, obj.args)
            self.update(h, obj.func)
            self.update(h, obj.keywords)
            return h.digest()

        elif type_util.is_pydantic_model(obj):
            try:
                # We have a choice to use pickle.dumps(), obj.model_dump(), or
                # obj.model_dump_json(). The advantage of pickle and JSON is that
                # we'd be treating Pydantic objects opaquely, so there are no surprises
                # for Pydantic users. And the reason to choose JSON over pickle is
                # that JSON is faster, weirdly enough.
                any_obj = cast("Any", obj)
                if hasattr(any_obj, "model_dump_json"):
                    # Pydantic v2
                    json_data = any_obj.model_dump_json()
                else:
                    # Pydantic v1
                    json_data = any_obj.json()
                self.update(h, json_data)
                return h.digest()
            except Exception as ex:
                raise UnhashableTypeError("""
                    Pydantic object contains unhashable members, such as functions.
                """) from ex

        else:
            # As a last resort, hash the output of the object's __reduce__ method
            try:
                reduce_data = obj.__reduce__()
            except Exception as ex:
                raise UnhashableTypeError() from ex

            for item in reduce_data:
                self.update(h, item)
            return h.digest()


class NoResult:
    """Placeholder class for return values when None is meaningful."""