import importlib.util
import sys
import types

import numpy as np


def _has_module(name):
    return importlib.util.find_spec(name) is not None


def _install_rocket_fft_stub():
    if _has_module("rocket_fft"):
        return
    sys.modules.setdefault("rocket_fft", types.ModuleType("rocket_fft"))


def _install_numba_stub():
    if _has_module("numba"):
        return

    numba_stub = types.ModuleType("numba")

    def njit(*args, **kwargs):
        # Supports both @njit and @njit(...)
        if args and callable(args[0]) and len(args) == 1 and not kwargs:
            return args[0]

        def decorator(func):
            return func

        return decorator

    numba_stub.njit = njit
    sys.modules.setdefault("numba", numba_stub)


def _install_pyfftw_stub():
    if _has_module("pyfftw"):
        return

    class FFTW:
        def __init__(self, input_array, output_array, direction="FFTW_FORWARD", flags=None):
            self._input = input_array
            self._output = output_array
            self._direction = direction

        def __call__(self):
            if self._direction == "FFTW_BACKWARD":
                self._output[:] = np.fft.ifft(self._input)
            else:
                self._output[:] = np.fft.rfft(self._input)

    pyfftw_stub = types.ModuleType("pyfftw")
    pyfftw_stub.empty_aligned = lambda size, dtype="float64": np.empty(size, dtype=dtype)
    pyfftw_stub.FFTW = FFTW
    sys.modules.setdefault("pyfftw", pyfftw_stub)


def _install_gwpy_stub():
    if _has_module("gwpy"):
        return

    gwpy_stub = types.ModuleType("gwpy")
    timeseries_stub = types.ModuleType("gwpy.timeseries")

    class TimeSeries:
        def __init__(self, data, sample_rate=None, t0=0):
            self.value = np.asarray(data)
            self.sample_rate = types.SimpleNamespace(value=sample_rate)
            self.t0 = types.SimpleNamespace(value=t0)

    timeseries_stub.TimeSeries = TimeSeries
    gwpy_stub.timeseries = timeseries_stub

    sys.modules.setdefault("gwpy", gwpy_stub)
    sys.modules.setdefault("gwpy.timeseries", timeseries_stub)


# Install stubs early during test collection so imports work in minimal environments.
_install_rocket_fft_stub()
_install_numba_stub()
_install_pyfftw_stub()
_install_gwpy_stub()
