import ctypes
import os
import platform

import numpy as np


class nfsft_plan(ctypes.Structure):
    pass


# NFSFT flags (aligned with pyNFFT3/flags.py)
NFSFT_NORMALIZED = 1 << 0
NFSFT_USE_NDFT = 1 << 1
NFSFT_USE_DPT = 1 << 2
NFSFT_MALLOC_X = 1 << 3
NFSFT_MALLOC_F_HAT = 1 << 5
NFSFT_MALLOC_F = 1 << 6
NFSFT_PRESERVE_F_HAT = 1 << 7
NFSFT_PRESERVE_X = 1 << 8
NFSFT_PRESERVE_F = 1 << 9
NFSFT_DESTROY_F_HAT = 1 << 10
NFSFT_DESTROY_X = 1 << 11
NFSFT_DESTROY_F = 1 << 12
NFSFT_NO_DIRECT_ALGORITHM = 1 << 13
NFSFT_NO_FAST_ALGORITHM = 1 << 14
NFSFT_ZERO_F_HAT = 1 << 16
NFSFT_EQUISPACED = 1 << 17

# NFFT flags used by NFSFT internals
PRE_PHI_HUT = 1 << 0
PRE_PSI = 1 << 4
FFTW_INIT = 1 << 10
MALLOC_X = 1 << 6
MALLOC_F_HAT = 1 << 7
MALLOC_F = 1 << 8
FFT_OUT_OF_PLACE = 1 << 9
NFFT_SORT_NODES = 1 << 11

# Defaults (aligned with pyNFFT3)
nfsft_default = NFSFT_MALLOC_X | NFSFT_MALLOC_F | NFSFT_MALLOC_F_HAT
nfsft_nfft_default = ctypes.c_uint32(
    PRE_PHI_HUT
    | PRE_PSI
    | FFTW_INIT
    | MALLOC_X
    | MALLOC_F_HAT
    | MALLOC_F
    | FFT_OUT_OF_PLACE
    | NFFT_SORT_NODES
)
nfsft_default_nfft_cut_off = 6


def _glibc_subdir() -> str:
    if os.name == "nt" or platform.system() == "Darwin":
        return ""

    libc = ctypes.CDLL(None)
    libc.gnu_get_libc_version.restype = ctypes.c_char_p
    version = libc.gnu_get_libc_version().decode("ascii")

    parts = version.split(".")
    major = int(parts[0])
    minor = int(parts[1]) if len(parts) > 1 else 0

    if (major, minor) < (2, 36):
        return "glibc2.22"
    return "glibc2.40"


def _lib_path() -> str:
    here = os.path.dirname(__file__)
    if os.name == "nt":
        return os.path.join(here, "libnfsftjulia.dll")
    if platform.system() == "Darwin":
        return os.path.join(here, "libnfsftjulia.dylib")
    return os.path.join(here, _glibc_subdir(), "libnfsftjulia.so")


def _set_signatures(lib):
    lib.jnfsft_alloc.restype = ctypes.POINTER(nfsft_plan)

    lib.jnfsft_init.argtypes = [
        ctypes.POINTER(nfsft_plan),
        ctypes.c_int32,
        ctypes.c_int32,
        ctypes.c_uint32,
        ctypes.c_uint32,
        ctypes.c_int32,
    ]

    lib.jnfsft_finalize.argtypes = [ctypes.POINTER(nfsft_plan)]

    lib.jnfsft_set_x.argtypes = [
        ctypes.POINTER(nfsft_plan),
        np.ctypeslib.ndpointer(np.float64, flags="C"),
    ]
    lib.jnfsft_set_x.restype = ctypes.POINTER(ctypes.c_double)

    lib.jnfsft_set_f.argtypes = [
        ctypes.POINTER(nfsft_plan),
        np.ctypeslib.ndpointer(np.complex128, flags="C"),
    ]
    lib.jnfsft_set_f.restype = ctypes.POINTER(ctypes.c_double)

    lib.jnfsft_set_fhat.argtypes = [
        ctypes.POINTER(nfsft_plan),
        np.ctypeslib.ndpointer(np.complex128, flags="C"),
    ]
    lib.jnfsft_set_fhat.restype = ctypes.POINTER(ctypes.c_double)

    lib.jnfsft_trafo.argtypes = [ctypes.POINTER(nfsft_plan)]
    lib.jnfsft_trafo.restype = ctypes.POINTER(ctypes.c_double)

    lib.jnfsft_adjoint.argtypes = [ctypes.POINTER(nfsft_plan)]
    lib.jnfsft_adjoint.restype = ctypes.POINTER(ctypes.c_double)

    lib.jnfsft_trafo_direct.argtypes = [ctypes.POINTER(nfsft_plan)]
    lib.jnfsft_trafo_direct.restype = ctypes.POINTER(ctypes.c_double)

    lib.jnfsft_adjoint_direct.argtypes = [ctypes.POINTER(nfsft_plan)]
    lib.jnfsft_adjoint_direct.restype = ctypes.POINTER(ctypes.c_double)

    lib.nfft_get_num_threads.restype = ctypes.c_int


class NFSFT:
    def __init__(
        self,
        N,
        M,
        flags=nfsft_default,
        nfft_flags=nfsft_nfft_default,
        nfft_cutoff=nfsft_default_nfft_cut_off,
    ):
        self.N = int(N)
        self.M = int(M)
        if self.N <= 0 or self.M <= 0:
            raise ValueError("N and M must be positive")

        self.flags = int(flags.value) if hasattr(flags, "value") else int(flags)
        if (self.flags & NFSFT_EQUISPACED) != 0 and self.__class__ is NFSFT:
            raise ValueError("Use NSFT class for equispaced variant")

        self.nfft_flags = (
            int(nfft_flags.value) if hasattr(nfft_flags, "value") else int(nfft_flags)
        )
        self.nfft_cutoff = int(nfft_cutoff)

        self.lib_path = _lib_path()
        self.lib = ctypes.CDLL(self.lib_path)
        _set_signatures(self.lib)

        self._plan = self.lib.jnfsft_alloc()
        self.lib.jnfsft_init(
            self._plan,
            ctypes.c_int32(self.N),
            ctypes.c_int32(self.M),
            ctypes.c_uint32(self.flags),
            ctypes.c_uint32(self.nfft_flags),
            ctypes.c_int32(self.nfft_cutoff),
        )

        self.finalized = False
        self._x = None
        self._f = None
        self._fhat = None

    def __del__(self):
        # Explicit finalize() is preferred; implicit finalize at interpreter shutdown
        # can trigger crashes in some libnfsft builds.
        return None

    @property
    def n_total(self):
        return (2 * self.N + 2) ** 2

    @property
    def num_threads(self):
        return int(self.lib.nfft_get_num_threads())

    def nfsft_index(self, k, n):
        return (2 * self.N + 2) * (self.N - n + 1) + (self.N + k + 1)

    @property
    def x(self):
        return self._x

    @x.setter
    def x(self, value):
        arr = np.asarray(value, dtype=np.float64, order="C")
        if arr.shape != (2, self.M):
            raise ValueError(f"x must have shape (2, {self.M})")
        ptr = self.lib.jnfsft_set_x(self._plan, np.ascontiguousarray(arr.T))
        self._x = np.ctypeslib.as_array(ptr, shape=(2 * self.M,)).reshape((2, self.M))

    @property
    def f(self):
        return self._f

    @f.setter
    def f(self, value):
        arr = np.asarray(value, dtype=np.complex128, order="C")
        if arr.shape != (self.M,):
            raise ValueError(f"f must have shape ({self.M},)")
        ptr = self.lib.jnfsft_set_f(self._plan, arr)
        self._f = np.ctypeslib.as_array(ptr, shape=(2 * self.M,)).view(np.complex128)

    @property
    def fhat(self):
        return self._fhat

    @fhat.setter
    def fhat(self, value):
        arr = np.asarray(value, dtype=np.complex128, order="C")
        if arr.shape != (self.n_total,):
            raise ValueError(f"fhat must have shape ({self.n_total},)")
        ptr = self.lib.jnfsft_set_fhat(self._plan, arr)
        self._fhat = np.ctypeslib.as_array(ptr, shape=(2 * self.n_total,)).view(np.complex128)

    def trafo(self):
        ptr = self.lib.jnfsft_trafo(self._plan)
        self._f = np.ctypeslib.as_array(ptr, shape=(2 * self.M,)).view(np.complex128)
        return self._f

    def trafo_direct(self):
        ptr = self.lib.jnfsft_trafo_direct(self._plan)
        self._f = np.ctypeslib.as_array(ptr, shape=(2 * self.M,)).view(np.complex128)
        return self._f

    def adjoint(self):
        ptr = self.lib.jnfsft_adjoint(self._plan)
        self._fhat = np.ctypeslib.as_array(ptr, shape=(2 * self.n_total,)).view(np.complex128)
        return self._fhat

    def adjoint_direct(self):
        ptr = self.lib.jnfsft_adjoint_direct(self._plan)
        self._fhat = np.ctypeslib.as_array(ptr, shape=(2 * self.n_total,)).view(np.complex128)
        return self._fhat

    def finalize(self):
        if not self.finalized:
            self.lib.jnfsft_finalize(self._plan)
            self.finalized = True


class NSFT(NFSFT):
    """Equispaced spherical transform variant (FSFT in pyNFFT3 naming)."""

    def __init__(
        self,
        N,
        M,
        flags=(nfsft_default | NFSFT_EQUISPACED),
        nfft_flags=nfsft_nfft_default,
        nfft_cutoff=nfsft_default_nfft_cut_off,
    ):
        super().__init__(N, M, flags=flags, nfft_flags=nfft_flags, nfft_cutoff=nfft_cutoff)

    @NFSFT.x.setter
    def x(self, value):
        raise RuntimeError("NSFT uses equispaced nodes internally; x cannot be set.")
