### import overloads, followed with dynamic implementation of functions
### we use separate file to avoid importing types
### overloads are generated by tools/generate_jax_random_overloads.py
### and are essential for static type analysis

from jax.random import * # first define all functions as in random
from jaxpy._random_jax_overloads import * # then overload with our definitions

import jax.random as _jax_random

import inspect as _inspect

# FIXME this needs to be acessible here, and singleton...
# FIXME we should be able to set the seed, and get the next key from jaxpy main
# jp._global_rng.__init__(0) would work so just make functions
from flax.nnx import Rngs as _Rngs
_global_rng = _Rngs(0)

# FIXME this could be usefull for front end user
# although for frontend we would just like to add __jaxpy_random_function__ to the function
# and maybe check if a key=None exists? I'm actually not sure
# another approach is just to do this, expect first argument to be key
def _jaxpy_random_wrapper_function(f, __module__ = "jaxpy.random_jax"):
    import functools
    # we need to remove the first parameter (key) from the signature
    sig = _inspect.signature(f)
    params = list(sig.parameters.values())

    if not params or params[0].name != 'key':
        raise ValueError(f"Function {f.__name__} does not have 'key' as first parameter.")

    # Remove 'key' from first position
    key_param = params.pop(0)
    # Re-add it as a keyword-only parameter with default=None
    key_param = key_param.replace(default=None, kind=_inspect.Parameter.KEYWORD_ONLY)
    params.append(key_param)

    new_sig = sig.replace(parameters=params)

    @functools.wraps(f)
    def wrapped(*args, key=None, **kwargs):
        if key is None:
            key = _global_rng.get_next()

        return f(key, *args, **kwargs)

    wrapped.__jaxpy_random_function__ = True
    wrapped.__signature__ = new_sig
    wrapped.__module__ = __module__ # set function module, could not make it work automatically, thus manual

    return wrapped






# all needs to be static for pylance to work properly
__all__ = []

# Iterate over all attributes in the module
for _name in dir(_jax_random):
    _attr = getattr(_jax_random, _name)
    if callable(_attr):
      # if first argument is key, we wrap the function with automatic key handling
      _sig = _inspect.signature(_attr)
      if len(_sig.parameters) > 0 and list(_sig.parameters.keys())[0] == 'key':
          _attr = _jaxpy_random_wrapper_function(_attr)
    # inject into global scope
    globals()[_name] = _attr
    __all__.append(_name)

