"""Utility to convert a user provided function into a Runnable with a ChannelWrite."""

from __future__ import annotations

import concurrent.futures
import functools
import inspect
import sys
import types
from collections.abc import Awaitable, Callable, Generator, Sequence
from datetime import timedelta
from typing import Any, Generic, TypeVar, cast

from langchain_core.runnables import Runnable
from typing_extensions import ParamSpec

from langgraph._internal._constants import CONF, CONFIG_KEY_CALL, RETURN
from langgraph._internal._runnable import (
    RunnableCallable,
    RunnableSeq,
    is_async_callable,
    run_in_executor,
)
from langgraph._internal._timeout import (
    coerce_timeout_policy,
    sync_timeout_unsupported,
)
from langgraph.config import get_config
from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry
from langgraph.types import CachePolicy, RetryPolicy, TimeoutPolicy

##
# Utilities borrowed from cloudpickle.
# https://github.com/cloudpipe/cloudpickle/blob/6220b0ce83ffee5e47e06770a1ee38ca9e47c850/cloudpickle/cloudpickle.py#L265


def _getattribute(obj: Any, name: str) -> Any:
    parent = None
    for subpath in name.split("."):
        if subpath == "<locals>":
            raise AttributeError(f"Can't get local attribute {name!r} on {obj!r}")
        try:
            parent = obj
            obj = getattr(obj, subpath)
        except AttributeError:
            raise AttributeError(f"Can't get attribute {name!r} on {obj!r}") from None
    return obj, parent


def _whichmodule(obj: Any, name: str) -> str | None:
    """Find the module an object belongs to.

    This function differs from ``pickle.whichmodule`` in two ways:
    - it does not mangle the cases where obj's module is __main__ and obj was
      not found in any module.
    - Errors arising during module introspection are ignored, as those errors
      are considered unwanted side effects.
    """
    module_name = getattr(obj, "__module__", None)

    if module_name is not None:
        return module_name
    # Protect the iteration by using a copy of sys.modules against dynamic
    # modules that trigger imports of other modules upon calls to getattr or
    # other threads importing at the same time.
    for module_name, module in sys.modules.copy().items():
        # Some modules such as coverage can inject non-module objects inside
        # sys.modules
        if (
            module_name == "__main__"
            or module_name == "__mp_main__"
            or module is None
            or not isinstance(module, types.ModuleType)
        ):
            continue
        try:
            if _getattribute(module, name)[0] is obj:
                return module_name
        except Exception:
            pass
    return None


def identifier(obj: Any, name: str | None = None) -> str | None:
    """Return the module and name of an object."""
    from langgraph._internal._runnable import RunnableCallable, RunnableSeq
    from langgraph.pregel._read import PregelNode

    if isinstance(obj, PregelNode):
        obj = obj.bound
    if isinstance(obj, RunnableSeq):
        obj = obj.steps[0]
    if isinstance(obj, RunnableCallable):
        obj = obj.func
    if name is None:
        name = getattr(obj, "__qualname__", None)
    if name is None:  # pragma: no cover
        # This used to be needed for Python 2.7 support but is probably not
        # needed anymore. However we keep the __name__ introspection in case
        # users of cloudpickle rely on this old behavior for unknown reasons.
        name = getattr(obj, "__name__", None)
    if name is None:
        return None

    module_name = getattr(obj, "__module__", None)
    if module_name is None:
        # In this case, obj.__module__ is None. obj is thus treated as dynamic.
        return None

    return f"{module_name}.{name}"


def _lookup_module_and_qualname(
    obj: Any, name: str | None = None
) -> tuple[types.ModuleType, str] | None:
    if name is None:
        name = getattr(obj, "__qualname__", None)
    if name is None:  # pragma: no cover
        # This used to be needed for Python 2.7 support but is probably not
        # needed anymore. However we keep the __name__ introspection in case
        # users of cloudpickle rely on this old behavior for unknown reasons.
        name = getattr(obj, "__name__", None)
    if name is None:
        return None

    module_name = _whichmodule(obj, name)

    if module_name is None:
        # In this case, obj.__module__ is None AND obj was not found in any
        # imported module. obj is thus treated as dynamic.
        return None

    if module_name == "__main__":
        return None

    # Note: if module_name is in sys.modules, the corresponding module is
    # assumed importable at unpickling time. See #357
    module = sys.modules.get(module_name, None)
    if module is None:
        # The main reason why obj's module would not be imported is that this
        # module has been dynamically created, using for example
        # types.ModuleType. The other possibility is that module was removed
        # from sys.modules after obj was created/imported. But this case is not
        # supported, as the standard pickle does not support it either.
        return None

    try:
        obj2, parent = _getattribute(module, name)
    except AttributeError:
        # obj was not found inside the module it points to
        return None
    if obj2 is not obj:
        return None
    return module, name


def _explode_args_trace_inputs(
    sig: inspect.Signature, input: tuple[tuple[Any, ...], dict[str, Any]]
) -> dict[str, Any]:
    args, kwargs = input
    bound = sig.bind_partial(*args, **kwargs)
    bound.apply_defaults()
    arguments = dict(bound.arguments)
    arguments.pop("self", None)
    arguments.pop("cls", None)
    for param_name, param in sig.parameters.items():
        if param.kind == inspect.Parameter.VAR_KEYWORD:
            # Update with the **kwargs, and remove the original entry
            # This is to help flatten out keyword arguments
            if param_name in arguments:
                arguments.update(arguments.pop(param_name))
    return arguments


def get_runnable_for_entrypoint(func: Callable[..., Any]) -> Runnable:
    key = (func, False)
    if key in CACHE:
        return CACHE[key]
    else:
        if is_async_callable(func):
            run = RunnableCallable(
                None, func, name=func.__name__, trace=False, recurse=False
            )
        else:
            afunc = functools.update_wrapper(
                functools.partial(run_in_executor, None, func), func
            )
            run = RunnableCallable(
                func,
                afunc,
                name=func.__name__,
                trace=False,
                recurse=False,
            )
        if not _lookup_module_and_qualname(func):
            return run
        return CACHE.setdefault(key, run)


def get_runnable_for_task(func: Callable[..., Any]) -> Runnable:
    key = (func, True)
    if key in CACHE:
        return CACHE[key]
    else:
        if hasattr(func, "__name__"):
            name = func.__name__
        elif hasattr(func, "func"):
            name = func.func.__name__
        elif hasattr(func, "__class__"):
            name = func.__class__.__name__
        else:
            name = str(func)

        if is_async_callable(func):
            run = RunnableCallable(
                None,
                func,
                explode_args=True,
                name=name,
                trace=False,
                recurse=False,
            )
        else:
            run = RunnableCallable(
                func,
                functools.wraps(func)(functools.partial(run_in_executor, None, func)),
                explode_args=True,
                name=name,
                trace=False,
                recurse=False,
            )
        seq = RunnableSeq(
            run,
            ChannelWrite([ChannelWriteEntry(RETURN)]),
            name=name,
            trace_inputs=functools.partial(
                _explode_args_trace_inputs, inspect.signature(func)
            ),
        )
        if not _lookup_module_and_qualname(func):
            return seq
        return CACHE.setdefault(key, seq)


CACHE: dict[tuple[Callable[..., Any], bool], Runnable] = {}


P = ParamSpec("P")
P1 = TypeVar("P1")
T = TypeVar("T")


class SyncAsyncFuture(Generic[T], concurrent.futures.Future[T]):
    def __await__(self) -> Generator[T, None, T]:
        yield cast(T, ...)


def call(
    func: Callable[P, Awaitable[T]] | Callable[P, T],
    *args: Any,
    retry_policy: Sequence[RetryPolicy] | None = None,
    cache_policy: CachePolicy | None = None,
    timeout: float | timedelta | TimeoutPolicy | None = None,
    **kwargs: Any,
) -> SyncAsyncFuture[T]:
    return _call_with_options(
        func,
        args,
        kwargs,
        retry_policy=retry_policy,
        cache_policy=cache_policy,
        timeout=coerce_timeout_policy(timeout),
    )


def _call_with_options(
    func: Callable[P, Awaitable[T]] | Callable[P, T],
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
    *,
    retry_policy: Sequence[RetryPolicy] | None = None,
    cache_policy: CachePolicy | None = None,
    timeout: TimeoutPolicy | None = None,
) -> SyncAsyncFuture[T]:
    if timeout is not None and not is_async_callable(func):
        name = getattr(func, "__name__", func.__class__.__name__)
        raise sync_timeout_unsupported(name, kind="Task")
    config = get_config()
    impl = config[CONF][CONFIG_KEY_CALL]
    fut = impl(
        func,
        (args, kwargs),
        retry_policy=retry_policy,
        cache_policy=cache_policy,
        callbacks=config["callbacks"],
        timeout=timeout,
    )
    return fut