# Copyright 2026 Marimo. All rights reserved.
from __future__ import annotations

import inspect
import numbers
import weakref
from collections.abc import Callable
from typing import TYPE_CHECKING, Any

from marimo._ast.cell import Cell
from marimo._ast.visitor import Name, VariableData

if TYPE_CHECKING:
    from marimo._runtime.dataflow import DirectedGraph

PRIMITIVES: tuple[type, ...] = (bytes, str, numbers.Number, type(None))
# Weakref instances should be disassociated from related references, as should
# other "primitives" as they are results and hopefully not hiding some scoped
# reference.
CLONE_PRIMITIVES = (weakref.ref,) + PRIMITIVES

FN_CACHE_TYPE = dict[Callable[..., Any] | type, bool] | None

UNCLONABLE_TYPES = [
    "marimo._runtime.state.State",
    "marimo._runtime.state.SetFunctor",
    "marimo._runtime.watch._file.FileState",
    "marimo._runtime.watch._directory.DirectoryState",
]

UNCLONABLE_MODULES = {
    "_asyncio",
    "_io",
    "marimo._ast",
    "marimo._plugins.ui",
    "numpy.lib.npyio",
}


def is_external(value: Any) -> bool:
    return "_marimo__cell_" not in inspect.getfile(value)


def is_primitive(value: Any) -> bool:
    # Tuples don't allow for write access
    if isinstance(value, tuple):
        return all(map(is_primitive, value))
    return isinstance(value, PRIMITIVES)


def is_primitive_type(value: type) -> bool:
    return any(issubclass(value, primitive) for primitive in PRIMITIVES)


def is_clone_primitive(value: Any) -> bool:
    return isinstance(value, CLONE_PRIMITIVES)


def is_data_primitive(value: Any) -> bool:
    if is_primitive(value):
        return True

    if inspect.isclass(value):
        return False

    if not (
        hasattr(value, "__array__")
        or hasattr(value, "toarray")
        or hasattr(value, "__array_interface__")
    ):
        return False

    # Handle cross device tensors particular to torch
    # Transfer may be expensive, so non-cpu tensors are considered
    # unhashable.
    if is_instance_by_name(value, "torch.Tensor"):
        return str(value.device) == "cpu"

    # If a numpy like array, ensure that it's not an object array.
    if hasattr(value, "dtype"):
        return not (
            value.dtype is None
            or (hasattr(value.dtype, "hasobject") and value.dtype.hasobject)
        )
    elif hasattr(value, "dtypes"):
        # Bit of discrepancy between objects like polars and pandas, so use
        # narwhals to normalize the dataframe.
        import narwhals.stable.v2 as nw

        try:
            return bool(
                all(
                    dtype.is_numeric()
                    for dtype in nw.from_native(value)
                    .collect_schema()
                    .dtypes()
                )
            )
        except Exception:
            return False
    # Otherwise may be a closely related array object
    return True


def _is_primitive_container(
    value: Any, predicate: Callable[[Any], bool]
) -> bool:
    visited = set()

    def recurse_container(value: Any) -> bool:
        if is_primitive(value):
            return True

        if id(value) in visited:
            return True

        if isinstance(value, dict):
            visited.add(id(value))
            return all(map(predicate, value.items()))
        # Tuple has to be considered too, since a tuple can contain containers.
        if isinstance(value, (set, list, tuple)):
            visited.add(id(value))
            return all(map(predicate, value))

        return False

    return recurse_container(value)


def is_data_primitive_container(value: Any) -> bool:
    return _is_primitive_container(value, is_data_primitive)


def is_primitive_container(value: Any) -> bool:
    return _is_primitive_container(value, is_primitive)


def is_instance_by_name(obj: object, name: str) -> bool:
    if not (hasattr(obj, "__module__") and hasattr(obj, "__class__")):
        return False
    obj_name = f"{obj.__module__}.{obj.__class__.__name__}"
    return obj_name == name


def is_unclonable_type(obj: object) -> bool:
    # Cell objects in particular are hidden by functools.wraps.
    if isinstance(obj, Cell):
        return True
    return any(is_instance_by_name(obj, name) for name in UNCLONABLE_TYPES)


def from_unclonable_module(obj: object) -> bool:
    obj = obj if hasattr(obj, "__module__") else obj.__class__
    return hasattr(obj, "__module__") and any(
        obj.__module__.startswith(name) for name in UNCLONABLE_MODULES
    )


def is_pure_scope(
    ref: Name,
    defs: dict[str, Any],
    cache: FN_CACHE_TYPE = None,
) -> bool:
    return inspect.ismodule(defs[ref]) or is_pure_function(
        ref, defs[ref], defs, cache
    )


def is_pure_function(
    ref: Name,
    value: Any,
    defs: dict[str, Any],
    cache: FN_CACHE_TYPE = None,
    graph: DirectedGraph | None = None,
) -> bool:
    if cache is None:
        cache = {}
    # Explicit removal of __hash__ indicates this is potentially mutable.
    # (e.g. list)
    if getattr(value, "__hash__", None) is None:
        return False
    try:
        if value in cache:
            return cache[value]
    except TypeError:
        # If value is not hashable, we assume it's not pure.
        return False
    # Trivial enough not to cache.
    if inspect.isclass(value) or not callable(value):
        return False
    # builtin_function_or_method are C functions, which we assume to be pure
    if inspect.isbuiltin(value):
        return True

    # Note: isfunction still covers lambdas, coroutines, etc.
    if not inspect.isfunction(value):
        return False

    # If this object wraps another, call the check on the wrapped object.
    if hasattr(value, "__wrapped__"):
        wrapped = getattr(value, "__wrapped__", None)
        if wrapped is not None:
            # This still catches impure decorated functions since the impure
            # reference will be captured by required refs.
            return is_pure_function(ref, wrapped, defs, cache, graph)

    # We assume all external module function references to be pure. Cache can
    # still be be invalidated by pin_modules attribute. Note this also captures
    # cases like functors from an external module.
    module = getattr(value, "__module__", None)
    is_lib = module != "__main__"
    # If the function exists on the graph as a top level function, it was
    # defined in the current notebook context. If is_lib is false, then it's a
    # marimo notebook as an external module, and we should still check for
    # 'purity'.
    is_marimo_lib = False
    if graph is not None:
        from marimo._ast.toplevel import TopLevelExtraction

        extraction = TopLevelExtraction.from_graph(graph)
        is_marimo_lib = ref in extraction.variables
    if is_lib and not is_marimo_lib:
        return True

    cache[value] = True  # Prevent recursion

    def cancel_predicate(ref: Name, _data: VariableData) -> bool:
        if not cache[value]:
            return False

        # A pure function can only refer to other functions, classes, or
        # modules.
        # External variable reference makes it inherently impure.
        if ref in defs:
            # Recursion allows for effective DFS
            if not (
                inspect.ismodule(defs[ref])
                or is_pure_function(ref, defs[ref], defs, cache, graph)
            ):
                cache[value] = False
                return False
        return True

    if graph is not None:
        graph.get_transitive_references(
            {ref}, inclusive=False, predicate=cancel_predicate
        )

    return cache[value]


def build_ref_predicate_for_primitives(
    glbls: dict[str, Any],
    primitives: tuple[type, ...] | None = None,
) -> Callable[[Name, VariableData], bool]:
    """
    Builds a predicate function to determine if a reference should be included

    Args:
        glbls: The global variables dictionary to base the predicate on
        primitives: A tuple of types that should be considered as base types
    Returns:
        A function that takes a variable name and associated data and
        returns True if its reference should be included in a reference search.

    All declared variables are tied together under the graph of required_refs.
    Strict execution gets the minimum graph of definitions for execution.
    Certain definitions, like lambdas, functions, and classes contain an
    executable body and require their `required_refs` to be scope (included in
    this graph). This function determines if a potential reference should be
    included in the graph based on its computed type. Consider:

    >>> def foo():
    ...     return bar()

    here `foo` is a function with `bar` as a reference in the execution body,
    so if `foo` is a reference, both `bar` and `foo` should be included in the
    graph, otherwise we'll get a NameError on `bar` if `foo` is called.
    Compare that to:

    >>> x = foo()

    if `x` is the only reference, should `foo` be included in the graph? It
    depends on the context, so we defer to the type of `x` which has already
    been computed at this point. If `x` is a known 'primitive' type, and thus
    does not have an executable body, we can exclude `foo` from the graph.
    However, `foo` may return a object or another function, which in turn may
    have references; so if x doesn't match the very low bar 'primitive', its
    `required_refs` are included in the graph.

    NB: The builtin `inspect.getclosurevars` exists, but it fails on some of
    these edgecases.

    NB: lambdas, as anonymous functions, do not have a name to refer to them-
    so visitor injects the dummy variable `_lambda` into the `required_refs` to
    denote their presence.
    """

    if primitives is None:
        primitives = PRIMITIVES

    def check_ref(ref: Name) -> bool:
        return ref in glbls and (
            inspect.isfunction(glbls[ref])
            or inspect.ismodule(glbls[ref])
            or inspect.isclass(glbls[ref])
            or callable(glbls[ref])
        )

    def only_scoped_refs(ref: Name, data: VariableData) -> bool:
        # TODO: Other common types could be added here, like numpy arrays that
        # are not dtype=object, etc.. that are known not to be dependent on the
        # functions that created them.

        # This errs on the side of including too much, but that's a better user
        # experience than not having definitions available.
        return (
            ref in glbls
            and not isinstance(glbls[ref], primitives)
            and (
                "_lambda" in data.required_refs
                or any(map(check_ref, data.required_refs | {ref}))
            )
        )

    return only_scoped_refs
