"""Funciones Numba de bajo nivel para operar sobre arrays de uint64 (limbs)."""

from __future__ import annotations

import numpy as np
from numba import njit, types, int64, uint64


@njit(int64(uint64))
def _u64_bit_length(x: np.uint64) -> np.int64:
    """Devuelve cuántos bits hacen falta para representar x.

    Usa búsqueda binaria en 6 pasos fijos en vez de recorrer bit a bit.
    Para x == 0 devuelve 0.
    """
    if x == 0:
        return np.int64(0)
    bl = np.int64(0)
    if x >= np.uint64(1) << np.uint64(32):
        bl += 32
        x >>= np.uint64(32)
    if x >= np.uint64(1) << np.uint64(16):
        bl += 16
        x >>= np.uint64(16)
    if x >= np.uint64(1) << np.uint64(8):
        bl += 8
        x >>= np.uint64(8)
    if x >= np.uint64(1) << np.uint64(4):
        bl += 4
        x >>= np.uint64(4)
    if x >= np.uint64(1) << np.uint64(2):
        bl += 2
        x >>= np.uint64(2)
    if x >= np.uint64(1) << np.uint64(1):
        bl += 1
    return bl + 1


@njit(int64(uint64[:]))
def _trim_nbits(limbs: np.ndarray) -> np.int64:
    """Devuelve cuántos bits se necesitan realmente para representar el valor.

    Recorre los limbs desde el más significativo hacia abajo, descarta los
    que valen cero y cuenta los bits útiles del último limb no nulo.
    Si todos son cero devuelve 0.
    """
    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(uint64[:, :](uint64[:], int64, int64))
def split_equal_parts_core(limbs: np.ndarray, nbits: int, parts: int) -> np.ndarray:
    """Corta el valor en ``parts`` trozos de igual tamaño en bits.

    Reparte los bits de menor a mayor peso entre las filas del resultado.
    ``nbits`` debe ser divisible por ``parts``.

    Devuelve un array 2D (parts x words_por_trozo) donde cada fila es un
    trozo en little-endian y la fila 0 corresponde a los bits más bajos.
    """
    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(types.Tuple((uint64[:], int64))(uint64[:], int64, int64, int64))
def extract_range_core(
    limbs: np.ndarray, nbits: int, start: int, width: int,
) -> tuple[np.ndarray, np.int64]:
    """Recorta ``width`` bits consecutivos a partir de la posición ``start``.

    Funciona como un slice a nivel de bits: copia el rango
    [start, start+width) y lo devuelve como un valor independiente.

    Devuelve (limbs_resultado, nbits_resultado) ya recortados, sin ceros
    de más al final. Si width es 0 devuelve un array vacío y nbits 0.
    """
    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(types.Tuple((uint64[:], int64))(uint64[:, :], int64))
def concat_parts_core(
    parts_2d: np.ndarray, chunk_bits: int,
) -> tuple[np.ndarray, np.int64]:
    """Reensambla trozos de bits en un solo valor (inversa de split).

    Toma un array 2D donde cada fila es un trozo de ``chunk_bits`` bits
    en little-endian (fila 0 = bits más bajos) y los pega uno tras otro.

    Devuelve (limbs_resultado, nbits_resultado) ya recortados.
    """
    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(types.Tuple((uint64[:], int64))(uint64[:], uint64[:]))
def bitwise_or_core(
    a: np.ndarray, b: np.ndarray,
) -> tuple[np.ndarray, np.int64]:
    """Aplica OR bit a bit entre dos valores representados como arrays LE.

    Si los arrays tienen distinta longitud, el más corto se trata como si
    tuviera ceros en las posiciones que le faltan.

    Devuelve (limbs_resultado, nbits_resultado) ya recortados.
    """
    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(types.Tuple((uint64[:], int64))(uint64[:], int64))
def shift_left_core(
    limbs: np.ndarray, k: int,
) -> tuple[np.ndarray, np.int64]:
    """Desplaza el valor k bits hacia la izquierda (equivale a multiplicar por 2^k).

    Si k es negativo o el array está vacío, devuelve una copia sin cambios.

    Devuelve (limbs_resultado, nbits_resultado) ya recortados.
    """
    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
