"""Kernels Numba reutilizables para operaciones sobre limbs uint64."""

from __future__ import annotations

import numpy as np
from numba import njit


@njit
def _u64_bit_length(x):
    """Bit-length de un uint64 (Numba-friendly)."""
    bl = 0
    while x != 0:
        x >>= 1
        bl += 1
    return bl


@njit
def _trim_nbits(limbs):
    """Calcula nbits canónico a partir de un array de limbs uint64."""
    n = limbs.size
    if n == 0:
        return np.int64(0)
    i = n - 1
    while i >= 0 and limbs[i] == 0:
        i -= 1
    if i < 0:
        return np.int64(0)
    return np.int64(i * 64 + _u64_bit_length(limbs[i]))


@njit
def split_equal_parts_core(limbs, nbits, parts):
    """Divide limbs en ``parts`` trozos iguales.

    Parameters
    ----------
    limbs : np.ndarray[uint64]
        Limbs LE del valor.
    nbits : int
        Bits significativos del valor.
    parts : int
        Número de partes (nbits debe ser divisible por parts).

    Returns
    -------
    out : np.ndarray[uint64] de forma (parts, words_per_part)
        Cada fila es un chunk en LE.  Fila 0 = chunk LSB.
    """
    chunk_bits = nbits // parts
    if chunk_bits == 0:
        return np.empty((parts, 0), dtype=np.uint64)

    words_per_part = (chunk_bits + 63) >> 6
    out = np.zeros((parts, words_per_part), dtype=np.uint64)

    for p in range(parts):
        start_bit = p * chunk_bits
        word0 = start_bit >> 6
        shift = start_bit & 63

        if shift == 0:
            for w in range(words_per_part):
                s = word0 + w
                if s < limbs.size:
                    out[p, w] = limbs[s]
        else:
            for w in range(words_per_part):
                s0 = word0 + w
                s1 = s0 + 1
                lo = limbs[s0] if s0 < limbs.size else np.uint64(0)
                hi = limbs[s1] if s1 < limbs.size else np.uint64(0)
                out[p, w] = (lo >> shift) | (hi << (64 - shift))

        rem = chunk_bits & 63
        if rem != 0:
            out[p, words_per_part - 1] &= (np.uint64(1) << rem) - np.uint64(1)

    return out


@njit
def extract_range_core(limbs, nbits, start, width):
    """Extrae ``width`` bits desde posición ``start``.

    Parameters
    ----------
    limbs : np.ndarray[uint64]
        Limbs LE.
    nbits : int
        Bits significativos.
    start : int
        Posición del bit LSB del rango.
    width : int
        Cantidad de bits a extraer.

    Returns
    -------
    out : np.ndarray[uint64]
        Limbs LE del valor extraído.
    out_nbits : int64
        Bits significativos del resultado (trimmed).
    """
    if width == 0:
        return np.empty(0, dtype=np.uint64), np.int64(0)

    n_words = (width + 63) >> 6
    out = np.zeros(n_words, dtype=np.uint64)

    word0 = start >> 6
    shift = start & 63

    if shift == 0:
        for w in range(n_words):
            s = word0 + w
            if s < limbs.size:
                out[w] = limbs[s]
    else:
        for w in range(n_words):
            s0 = word0 + w
            s1 = s0 + 1
            lo = limbs[s0] if s0 < limbs.size else np.uint64(0)
            hi = limbs[s1] if s1 < limbs.size else np.uint64(0)
            out[w] = (lo >> shift) | (hi << (64 - shift))

    rem = width & 63
    if rem != 0:
        out[n_words - 1] &= (np.uint64(1) << rem) - np.uint64(1)

    out_nbits = _trim_nbits(out)
    if out_nbits == 0:
        return np.empty(0, dtype=np.uint64), np.int64(0)
    need = (out_nbits + 63) >> 6
    return out[:need].copy(), out_nbits


@njit
def concat_parts_core(parts_2d, chunk_bits):
    """Concatena partes 2D de vuelta en un array de limbs.

    Parameters
    ----------
    parts_2d : np.ndarray[uint64] de forma (n_parts, words_per_part)
        Cada fila es un chunk en LE. Fila 0 = chunk LSB.
    chunk_bits : int
        Bits por cada parte.

    Returns
    -------
    limbs : np.ndarray[uint64]
        Limbs LE concatenados.
    nbits : int64
        Bits significativos del resultado.
    """
    n_parts = parts_2d.shape[0]
    total_bits = n_parts * chunk_bits
    if total_bits == 0:
        return np.empty(0, dtype=np.uint64), np.int64(0)

    total_words = (total_bits + 63) >> 6
    out = np.zeros(total_words, dtype=np.uint64)

    for p in range(n_parts):
        start_bit = p * chunk_bits
        word0 = start_bit >> 6
        shift = start_bit & 63
        words_per_part = parts_2d.shape[1]

        for w in range(words_per_part):
            val = parts_2d[p, w]
            dst = word0 + w
            if shift == 0:
                if dst < total_words:
                    out[dst] |= val
            else:
                if dst < total_words:
                    out[dst] |= val << shift
                dst1 = dst + 1
                if dst1 < total_words:
                    out[dst1] |= val >> (64 - shift)

    out_nbits = _trim_nbits(out)
    if out_nbits == 0:
        return np.empty(0, dtype=np.uint64), np.int64(0)
    need = (out_nbits + 63) >> 6
    return out[:need].copy(), out_nbits


@njit
def bitwise_or_core(a, b):
    """OR element-wise de dos arrays LE uint64.

    Parameters
    ----------
    a, b : np.ndarray[uint64]
        Limbs LE.

    Returns
    -------
    out : np.ndarray[uint64]
        Limbs LE del resultado, longitud max(len(a), len(b)).
    out_nbits : int64
        Bits significativos del resultado (trimmed).
    """
    na = a.size
    nb = b.size
    n = max(na, nb)
    if n == 0:
        return np.empty(0, dtype=np.uint64), np.int64(0)
    out = np.zeros(n, dtype=np.uint64)
    for i in range(n):
        va = a[i] if i < na else np.uint64(0)
        vb = b[i] if i < nb else np.uint64(0)
        out[i] = va | vb
    out_nbits = _trim_nbits(out)
    if out_nbits == 0:
        return np.empty(0, dtype=np.uint64), np.int64(0)
    need = (out_nbits + 63) >> 6
    return out[:need].copy(), out_nbits


@njit
def shift_left_core(limbs, k):
    """Shift left por k bits sobre array LE uint64.

    Parameters
    ----------
    limbs : np.ndarray[uint64]
        Limbs LE.
    k : int
        Cantidad de bits a desplazar a la izquierda.

    Returns
    -------
    out : np.ndarray[uint64]
        Limbs LE del resultado.
    out_nbits : int64
        Bits significativos del resultado (trimmed).
    """
    if limbs.size == 0 or k < 0:
        return limbs.copy(), _trim_nbits(limbs)

    word_shift = k >> 6
    bit_shift = k & 63
    n = limbs.size + word_shift + 1  # +1 for possible overflow
    out = np.zeros(n, dtype=np.uint64)

    if bit_shift == 0:
        for i in range(limbs.size):
            out[i + word_shift] = limbs[i]
    else:
        for i in range(limbs.size):
            out[i + word_shift] |= limbs[i] << bit_shift
            out[i + word_shift + 1] |= limbs[i] >> (64 - bit_shift)

    out_nbits = _trim_nbits(out)
    if out_nbits == 0:
        return np.empty(0, dtype=np.uint64), np.int64(0)
    need = (out_nbits + 63) >> 6
    return out[:need].copy(), out_nbits
