from typing import Callable
from functools import partial

import jax.numpy as jnp
import jax.extend.core as jexc
from jaxtyping import DTypeLike
import equinox as eqx
import quax

from ._core import (
    _make_default_autocast_config,
    _AutocastConfig,
    grampify,
    ungrampify,
)


class _AutocastWrapper(eqx.Module):
    fn: Callable
    config: _AutocastConfig = eqx.field(static=True)

    def __call__(self, first_arg=None, *rest_args, **kwargs):
        first_arg = grampify(first_arg, config=self.config)
        out = quax.quaxify(self.fn)(first_arg, *rest_args, **kwargs)
        return ungrampify(out)

    @property
    def __wrapped__(self):
        return self.fn


def autocast(
    fn: Callable = None,
    /,
    dtype: DTypeLike = jnp.bfloat16,
    config: dict[jexc.Primitive, DTypeLike] | None = None,
) -> Callable:
    if fn is None:
        return partial(autocast, dtype=dtype, config=config)

    config = config or _make_default_autocast_config(dtype)

    return eqx.module_update_wrapper(
        _AutocastWrapper(fn=fn, config=_AutocastConfig(config))
    )
