"""Deep recursive comparator for verifying behavioral equivalence."""

from __future__ import annotations

import _thread
import array
import ast
import builtins
import datetime
import decimal
import enum
import io
import itertools
import logging
import math
import re
import sqlite3
import threading
import types
import warnings
import weakref
import xml.etree.ElementTree as ET
from collections import ChainMap, OrderedDict, deque
from importlib.util import find_spec
from typing import Any

log = logging.getLogger(__name__)

HAS_NUMPY = find_spec("numpy") is not None
HAS_SQLALCHEMY = find_spec("sqlalchemy") is not None
HAS_SCIPY = find_spec("scipy") is not None
HAS_PANDAS = find_spec("pandas") is not None
HAS_PYRSISTENT = find_spec("pyrsistent") is not None
HAS_TORCH = find_spec("torch") is not None
HAS_JAX = find_spec("jax") is not None
HAS_XARRAY = find_spec("xarray") is not None
HAS_TENSORFLOW = find_spec("tensorflow") is not None
HAS_NUMBA = find_spec("numba") is not None
HAS_PYARROW = find_spec("pyarrow") is not None

if HAS_NUMPY:
    import numpy as np  # type: ignore[import-not-found]
if HAS_SCIPY:
    import scipy  # type: ignore[import-untyped]
if HAS_JAX:
    import jax  # type: ignore[import-not-found]
    import jax.numpy as jnp  # type: ignore[import-not-found]
if HAS_XARRAY:
    import xarray  # type: ignore[import-not-found]
if HAS_TENSORFLOW:
    import tensorflow as tf  # type: ignore[import-untyped]
if HAS_SQLALCHEMY:
    import sqlalchemy  # type: ignore[import-not-found]
if HAS_PYARROW:
    import pyarrow as pa  # type: ignore[import-not-found]
if HAS_PANDAS:
    import pandas as pd  # type: ignore[import-untyped]
if HAS_TORCH:
    import torch  # type: ignore[import-not-found]
if HAS_NUMBA:
    import numba  # type: ignore[import-not-found]  # noqa: I001
    from numba.core.dispatcher import Dispatcher  # type: ignore[import-not-found]
    from numba.typed import Dict as NumbaDict  # type: ignore[import-not-found]
    from numba.typed import List as NumbaList
if HAS_PYRSISTENT:
    import pyrsistent  # type: ignore[import-not-found]

# Pattern to match pytest temp directories
PYTEST_TEMP_PATH_PATTERN = re.compile(
    r"/tmp/pytest-of-[^/]+/pytest-\d+/"  # noqa: S108
)

# Pattern to match Python tempfile directories
PYTHON_TEMPFILE_PATTERN = re.compile(
    r"/tmp/tmp[a-zA-Z0-9_]+/"  # noqa: S108
)

DICT_KEYS_TYPE: type[Any] = type({}.keys())
DICT_VALUES_TYPE: type[Any] = type({}.values())
DICT_ITEMS_TYPE: type[Any] = type({}.items())

# Fast-path types that can be compared with == directly.
# Uses type identity (not isinstance) for O(1) dispatch.
_IDENTITY_EQ_TYPES: frozenset[type[Any]] = frozenset(
    {
        int,
        bool,
        complex,
        type(None),
        type(Ellipsis),
        decimal.Decimal,
        set,
        bytes,
        bytearray,
        memoryview,
        frozenset,
        type,
        range,
        slice,
        OrderedDict,
        types.GenericAlias,
    }
)

EQUALITY_TYPES = (
    int,
    bool,
    complex,
    type(None),
    type(Ellipsis),
    decimal.Decimal,
    set,
    bytes,
    bytearray,
    memoryview,
    frozenset,
    enum.Enum,
    type,
    range,
    slice,
    OrderedDict,
    types.GenericAlias,
    *(
        (_union_type,)
        if (_union_type := getattr(types, "UnionType", None))
        else ()
    ),
)


def normalize_temp_path(path: str) -> str:
    """Normalize temporary file paths by replacing session-specific components."""
    path = PYTEST_TEMP_PATH_PATTERN.sub(
        "/tmp/pytest-temp/",  # noqa: S108
        path,
    )
    return PYTHON_TEMPFILE_PATTERN.sub(
        "/tmp/python-temp/",  # noqa: S108
        path,
    )


def is_temp_path(s: str) -> bool:
    """Check if a string looks like a temp path."""
    return (
        PYTEST_TEMP_PATH_PATTERN.search(s) is not None
        or PYTHON_TEMPFILE_PATTERN.search(s) is not None
    )


def extract_exception_from_message(
    msg: str,
) -> BaseException | None:
    """Try to extract a wrapped exception type from an error message."""
    match = re.search(r"got (\w+)\(['\"]", msg)
    if match:
        exc_name = match.group(1)
        exc_class = getattr(builtins, exc_name, None)
        if (
            exc_class is not None
            and isinstance(exc_class, type)
            and issubclass(exc_class, BaseException)
        ):
            result: BaseException = exc_class()
            return result
    return None


def get_wrapped_exception(
    exc: BaseException,
) -> BaseException | None:
    """Get the wrapped exception if this is a simple wrapper."""
    if hasattr(exc, "exceptions"):
        exceptions: Any = exc.exceptions
        if len(exceptions) == 1:
            inner: BaseException = exceptions[0]
            return inner
    if exc.__cause__ is not None:
        return exc.__cause__
    return extract_exception_from_message(str(exc))


def comparator(
    orig: Any,
    new: Any,
    superset_obj: bool = False,  # noqa: FBT001, FBT002
) -> bool:
    """
    Compare two objects for equality recursively.

    If *superset_obj* is True, the new object is allowed to have more
    keys than the original object.
    """
    try:
        return bool(_comparator_inner(orig, new, superset_obj))
    except RecursionError:
        log.exception("RecursionError while comparing objects")
        return False
    except Exception:
        log.exception("Error while comparing objects")
        return False


def _is_pickle_placeholder_error(exc: BaseException) -> bool:
    """Check if *exc* is a PicklePlaceholderAccessError by class name.

    We check by name to avoid importing the class at module level,
    which would pull in the picklepatch package.
    """
    return type(exc).__name__ == "PicklePlaceholderAccessError"


def _comparator_inner(  # noqa: C901, PLR0911, PLR0912, PLR0915
    orig: Any,
    new: Any,
    superset_obj: bool = False,  # noqa: FBT001, FBT002
) -> Any:
    """Recursively compare two values for deep behavioral equality."""
    # Handle exceptions specially
    if isinstance(orig, BaseException) and isinstance(new, BaseException):
        if _is_pickle_placeholder_error(
            orig,
        ) or _is_pickle_placeholder_error(new):
            log.debug(
                "Unable to verify behavior of unpickleable object "
                "in replay test",
            )
            return False
        if type(orig) is type(new):
            orig_dict = {
                k: v for k, v in orig.__dict__.items() if not k.startswith("_")
            }
            new_dict = {
                k: v for k, v in new.__dict__.items() if not k.startswith("_")
            }
            return comparator(orig_dict, new_dict, superset_obj)

        wrapped_orig = get_wrapped_exception(orig)
        if wrapped_orig is not None and comparator(
            wrapped_orig, new, superset_obj
        ):
            return True

        wrapped_new = get_wrapped_exception(new)
        return wrapped_new is not None and comparator(
            orig, wrapped_new, superset_obj
        )

    orig_type = type(orig)
    if type(orig) is not type(new):
        type_obj = orig_type
        new_type_obj = type(new)
        if (
            type_obj.__name__ != new_type_obj.__name__
            or type_obj.__qualname__ != new_type_obj.__qualname__
        ):
            return False

    # Fast-path: O(1) dispatch for the most common built-in types.
    # Uses type identity (not isinstance) to avoid walking the MRO.
    if orig_type is str:
        if orig == new:
            return True
        if is_temp_path(orig) and is_temp_path(new):
            return normalize_temp_path(orig) == normalize_temp_path(new)
        return False
    if orig_type is list or orig_type is tuple:
        if len(orig) != len(new):
            return False
        return all(
            comparator(elem1, elem2, superset_obj)
            for elem1, elem2 in zip(orig, new)
        )
    if orig_type is dict:
        if superset_obj:
            return all(
                k in new and comparator(v, new[k], superset_obj)
                for k, v in orig.items()
            )
        if len(orig) != len(new):
            return False
        for key in orig:
            if key not in new:
                return False
            if not comparator(orig[key], new[key], superset_obj):
                return False
        return True
    if orig_type is float:
        if math.isnan(orig) and math.isnan(new):
            return True
        return math.isclose(orig, new)
    if orig_type in _IDENTITY_EQ_TYPES:
        return orig == new

    # Slower isinstance path for subclasses and less common types
    if isinstance(orig, (list, tuple, deque, ChainMap)):
        if len(orig) != len(new):
            return False
        return all(
            comparator(elem1, elem2, superset_obj)
            for elem1, elem2 in zip(orig, new)
        )

    if isinstance(orig, str):
        if orig == new:
            return True
        if is_temp_path(orig) and is_temp_path(new):
            return normalize_temp_path(orig) == normalize_temp_path(new)
        return False

    if isinstance(orig, EQUALITY_TYPES):
        return orig == new
    if isinstance(orig, float):
        if math.isnan(orig) and math.isnan(new):
            return True
        return math.isclose(orig, new)

    if isinstance(orig, weakref.ref):
        orig_referent = orig()
        new_referent = new()
        if orig_referent is None and new_referent is None:
            return True
        if orig_referent is None or new_referent is None:
            return False
        return comparator(orig_referent, new_referent, superset_obj)

    if HAS_JAX and isinstance(orig, jax.Array):
        if orig.dtype != new.dtype:
            return False
        if orig.shape != new.shape:
            return False
        return bool(jnp.allclose(orig, new, equal_nan=True))

    if HAS_XARRAY and isinstance(orig, (xarray.Dataset, xarray.DataArray)):
        return orig.identical(new)

    if HAS_TENSORFLOW:
        if isinstance(orig, tf.Tensor):
            if orig.dtype != new.dtype:
                return False
            if orig.shape != new.shape:
                return False
            return comparator(orig.numpy(), new.numpy(), superset_obj)

        if isinstance(orig, tf.Variable):
            if orig.dtype != new.dtype:
                return False
            if orig.shape != new.shape:
                return False
            return comparator(orig.numpy(), new.numpy(), superset_obj)

        if isinstance(orig, tf.dtypes.DType):
            return orig == new

        if isinstance(orig, tf.TensorShape):
            return orig == new

        if isinstance(orig, tf.SparseTensor):
            if not comparator(
                orig.dense_shape.numpy(),
                new.dense_shape.numpy(),
                superset_obj,
            ):
                return False
            return comparator(
                orig.indices.numpy(),
                new.indices.numpy(),
                superset_obj,
            ) and comparator(
                orig.values.numpy(),
                new.values.numpy(),
                superset_obj,
            )

        if isinstance(orig, tf.RaggedTensor):
            if orig.dtype != new.dtype:
                return False
            if orig.shape.rank != new.shape.rank:
                return False
            return comparator(orig.to_list(), new.to_list(), superset_obj)

    if HAS_SQLALCHEMY:
        try:
            sqlalchemy.inspection.inspect(orig)
            sqlalchemy.inspection.inspect(new)
        except sqlalchemy.exc.NoInspectionAvailable:
            pass
        else:
            orig_keys = orig.__dict__
            new_keys = new.__dict__
            for key in list(orig_keys.keys()):
                if key.startswith("_"):
                    continue
                if key not in new_keys or not comparator(
                    orig_keys[key], new_keys[key], superset_obj
                ):
                    return False
            return True

    if isinstance(orig, dict) and not (
        HAS_SCIPY and isinstance(orig, scipy.sparse.spmatrix)
    ):
        if superset_obj:
            return all(
                k in new and comparator(v, new[k], superset_obj)
                for k, v in orig.items()
            )
        if len(orig) != len(new):
            return False
        for key in orig:
            if key not in new:
                return False
            if not comparator(orig[key], new[key], superset_obj):
                return False
        return True

    if isinstance(orig, types.MappingProxyType):
        return comparator(dict(orig), dict(new), superset_obj)

    if isinstance(orig, DICT_KEYS_TYPE):
        return comparator(set(orig), set(new))
    if isinstance(orig, DICT_VALUES_TYPE):
        return comparator(list(orig), list(new))
    if isinstance(orig, DICT_ITEMS_TYPE):
        return comparator(dict(orig), dict(new), superset_obj)

    if HAS_NUMPY:
        if isinstance(orig, (np.datetime64, np.timedelta64)):
            if np.isnat(orig) and np.isnat(new):
                return True
            if np.isnat(orig) or np.isnat(new):
                return False
            return orig == new

        if isinstance(orig, np.ndarray):
            if orig.dtype != new.dtype:
                return False
            if orig.shape != new.shape:
                return False
            if orig.ndim == 0:
                try:
                    return np.allclose(orig, new, equal_nan=True)
                except Exception:
                    return bool(orig == new)
            try:
                return np.allclose(orig, new, equal_nan=True)
            except Exception:
                return np.all(
                    [comparator(x, y, superset_obj) for x, y in zip(orig, new)]
                )

        if isinstance(orig, (np.floating, np.complexfloating)):
            return np.isclose(orig, new, equal_nan=True)

        if isinstance(orig, (np.integer, np.bool_, np.byte)):
            return orig == new

        if isinstance(orig, np.void):
            if orig.dtype != new.dtype:
                return False
            return all(
                comparator(orig[field], new[field], superset_obj)
                for field in orig.dtype.fields
            )

        if isinstance(orig, np.dtype):
            return orig == new

        if isinstance(orig, np.random.Generator):
            orig_state = orig.bit_generator.state
            new_state = new.bit_generator.state
            return comparator(orig_state, new_state, superset_obj)

        if isinstance(orig, np.random.RandomState):
            orig_state = orig.get_state(legacy=False)
            new_state = new.get_state(legacy=False)
            return comparator(orig_state, new_state, superset_obj)

    if HAS_SCIPY and isinstance(orig, scipy.sparse.spmatrix):
        if orig.dtype != new.dtype:
            return False
        if orig.get_shape() != new.get_shape():
            return False
        return (orig != new).nnz == 0

    if HAS_PYARROW:
        if isinstance(orig, pa.Table):
            if orig.schema != new.schema:
                return False
            if orig.num_rows != new.num_rows:
                return False
            return bool(orig.equals(new))

        if isinstance(orig, pa.RecordBatch):
            if orig.schema != new.schema:
                return False
            if orig.num_rows != new.num_rows:
                return False
            return bool(orig.equals(new))

        if isinstance(orig, pa.ChunkedArray):
            if orig.type != new.type:
                return False
            if len(orig) != len(new):
                return False
            return bool(orig.equals(new))

        if isinstance(orig, pa.Array):
            if orig.type != new.type:
                return False
            if len(orig) != len(new):
                return False
            return bool(orig.equals(new))

        if isinstance(orig, pa.Scalar):
            if orig.type != new.type:
                return False
            if not orig.is_valid and not new.is_valid:
                return True
            if not orig.is_valid or not new.is_valid:
                return False
            return bool(orig.equals(new))

        if isinstance(orig, (pa.Schema, pa.Field, pa.DataType)):
            return bool(orig.equals(new))

    if HAS_PANDAS:
        if isinstance(
            orig,
            (
                pd.DataFrame,
                pd.Series,
                pd.Index,
                pd.Categorical,
                pd.arrays.SparseArray,
            ),
        ):
            return bool(orig.equals(new))

        if isinstance(
            orig,
            (
                pd.CategoricalDtype,
                pd.Interval,
                pd.Period,
            ),
        ):
            return orig == new
        if pd.isna(orig) and pd.isna(new):
            return True

    if isinstance(orig, array.array):
        if orig.typecode != new.typecode:
            return False
        if len(orig) != len(new):
            return False
        return all(
            comparator(elem1, elem2, superset_obj)
            for elem1, elem2 in zip(orig, new)
        )

    try:
        if HAS_NUMPY and np.isnan(orig):
            return np.isnan(new)
    except Exception:  # noqa: S110
        pass
    try:
        if HAS_NUMPY and np.isinf(orig):
            return np.isinf(new)
    except Exception:  # noqa: S110
        pass

    if HAS_TORCH:
        if isinstance(orig, torch.Tensor):
            if orig.dtype != new.dtype:
                return False
            if orig.shape != new.shape:
                return False
            if orig.requires_grad != new.requires_grad:
                return False
            if orig.device != new.device:
                return False
            return torch.allclose(orig, new, equal_nan=True)

        if isinstance(orig, torch.dtype):
            return orig == new

        if isinstance(orig, torch.device):
            return orig == new

    if HAS_NUMBA:
        if isinstance(orig, NumbaList):
            if len(orig) != len(new):
                return False
            return all(
                comparator(elem1, elem2, superset_obj)
                for elem1, elem2 in zip(orig, new)
            )

        if isinstance(orig, NumbaDict):
            if superset_obj:
                return all(
                    key in new
                    and comparator(orig[key], new[key], superset_obj)
                    for key in orig
                )
            if len(orig) != len(new):
                return False
            for key in orig:
                if key not in new:
                    return False
                if not comparator(orig[key], new[key], superset_obj):
                    return False
            return True

        if isinstance(orig, numba.core.types.Type):
            return orig == new

        if isinstance(orig, Dispatcher):
            return orig.py_func is new.py_func

    if HAS_PYRSISTENT and isinstance(
        orig,
        (
            pyrsistent.PMap,
            pyrsistent.PVector,
            pyrsistent.PSet,
            pyrsistent.PRecord,
            pyrsistent.PClass,
            pyrsistent.PBag,
            pyrsistent.PList,
            pyrsistent.PDeque,
        ),
    ):
        return orig == new

    if hasattr(orig, "__attrs_attrs__") and hasattr(new, "__attrs_attrs__"):
        orig_dict = {}
        new_dict = {}

        for attr in orig.__attrs_attrs__:
            if attr.eq:
                attr_name = attr.name
                orig_dict[attr_name] = getattr(orig, attr_name, None)
                new_dict[attr_name] = getattr(new, attr_name, None)

        if superset_obj:
            new_attrs_dict = {}
            for attr in new.__attrs_attrs__:
                if attr.eq:
                    attr_name = attr.name
                    new_attrs_dict[attr_name] = getattr(new, attr_name, None)
            return all(
                k in new_attrs_dict
                and comparator(v, new_attrs_dict[k], superset_obj)
                for k, v in orig_dict.items()
            )
        return comparator(orig_dict, new_dict, superset_obj)

    if isinstance(orig, itertools.count):
        return repr(orig) == repr(new)

    if isinstance(orig, itertools.repeat):
        return repr(orig) == repr(new)

    if isinstance(orig, itertools.cycle):
        try:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", DeprecationWarning)
                orig_reduce = orig.__reduce__()
                new_reduce = new.__reduce__()
            orig_remaining = list(orig_reduce[1][0])
            new_remaining = list(new_reduce[1][0])
            orig_saved, orig_started = orig_reduce[2]  # type: ignore[str-unpack]
            new_saved, new_started = new_reduce[2]
            if orig_started != new_started:
                return False
            return comparator(
                orig_remaining, new_remaining, superset_obj
            ) and comparator(orig_saved, new_saved, superset_obj)
        except TypeError:
            sample_size = 200
            orig_sample = [next(orig) for _ in range(sample_size)]
            new_sample = [next(new) for _ in range(sample_size)]
            return comparator(orig_sample, new_sample, superset_obj)

    if type(orig).__module__ == "itertools":
        if isinstance(orig, itertools.groupby):
            orig_groups = [(k, list(g)) for k, g in orig]
            new_groups = [(k, list(g)) for k, g in new]
            return comparator(orig_groups, new_groups, superset_obj)
        return comparator(list(orig), list(new), superset_obj)

    if isinstance(
        orig,
        (
            datetime.datetime,
            datetime.date,
            datetime.timedelta,
            datetime.time,
            datetime.timezone,
            re.Pattern,
        ),
    ):
        return orig == new

    try:
        if hasattr(orig, "__eq__") and isinstance(
            orig.__eq__, types.MethodType
        ):
            return orig == new
    except Exception:  # noqa: S110
        pass

    if hasattr(orig, "__dict__") and hasattr(new, "__dict__"):
        orig_keys = orig.__dict__
        new_keys = new.__dict__
        if (
            type(orig_keys) is types.MappingProxyType
            and type(new_keys) is types.MappingProxyType
        ):
            if orig != new:
                return False
            orig_keys = dict(orig_keys)
            new_keys = dict(new_keys)
            orig_keys = {
                k: v for k, v in orig_keys.items() if not k.startswith("__")
            }
            new_keys = {
                k: v for k, v in new_keys.items() if not k.startswith("__")
            }

        if superset_obj:
            return all(
                k in new_keys and comparator(v, new_keys[k], superset_obj)
                for k, v in orig_keys.items()
            )

        if isinstance(orig, ast.AST):
            orig_keys = {
                k: v for k, v in orig.__dict__.items() if k != "parent"
            }
            new_keys = {k: v for k, v in new.__dict__.items() if k != "parent"}
        return comparator(orig_keys, new_keys, superset_obj)

    if hasattr(type(orig), "__slots__"):
        all_slots = set()
        for cls in type(orig).__mro__:
            if hasattr(cls, "__slots__"):
                all_slots.update(cls.__slots__)
        orig_vals = {s: getattr(orig, s, None) for s in all_slots}
        new_vals = {s: getattr(new, s, None) for s in all_slots}
        if superset_obj:
            return all(
                k in new_vals and comparator(v, new_vals[k], superset_obj)
                for k, v in orig_vals.items()
            )
        return comparator(orig_vals, new_vals, superset_obj)

    if type(orig) in {
        types.BuiltinFunctionType,
        types.BuiltinMethodType,
    }:
        return new == orig
    if isinstance(orig, ET.Element):
        return isinstance(new, ET.Element) and ET.tostring(
            orig
        ) == ET.tostring(new)
    if isinstance(
        orig,
        (
            _thread.LockType,
            _thread.RLock,
            threading.Event,
            threading.Condition,
            sqlite3.Connection,
            sqlite3.Cursor,
            io.IOBase,
        ),
    ):
        return type(orig) is type(new)
    if str(type(orig)) == "<class 'object'>":
        return True
    log.warning("Unknown comparator input type: %s", type(orig))
    return False
