import operator
import random
import pytest

# Ajusta este import a tu proyecto
from hugenat import HugeNat

import numpy as np
from numba import njit

WORD_BITS = 64


# ----------------------------
# Oráculos de referencia (int)
# ----------------------------

def ref_bit(x: int, i: int) -> int:
    """
    Oráculo de referencia para leer un bit de un entero Python.

    - Solo admite naturales (x >= 0). Si no, lanza ValueError.
    - Usa la misma convención que HugeNat:
      * i >= 0 indexa desde el bit menos significativo (LSB = bit 0).
      * i < 0 indexa desde el bit más significativo del número (relativo a bit_length).
    - Si el índice cae fuera del rango de bits representados (0..nbits-1), devuelve 0.

    Devuelve 0 o 1 como entero Python.
    """
    if x < 0:
        raise ValueError
    nbits = x.bit_length()
    if i < 0:
        if nbits == 0:
            raise IndexError
        i += nbits
    if i < 0 or i >= nbits:
        return 0
    return (x >> i) & 1


def ref_slice_step1(x: int, start, stop) -> int:
    """
    Oráculo de referencia para slicing rápido (step None/1) sobre un entero Python.

    Emula la semántica que queremos para el fast-path de HugeNat:
    - Interpreta `start` y `stop` como lo hace Python (incluye negativos relativos a bit_length).
    - Luego recorta el intervalo a [0, nbits] para evitar “bits infinitos a 0”.
    - Extrae los bits en el rango [start, stop) y los recompacta:
      el bit `start` pasa a ser el bit 0 del resultado.
    - Si el slice queda vacío (start >= stop), devuelve 0.

    Devuelve un entero Python con el valor extraído.
    """
    if x < 0:
        raise ValueError
    nbits = x.bit_length()

    if start is None:
        start = 0
    if stop is None:
        stop = nbits

    if start < 0:
        start += nbits
    if stop < 0:
        stop += nbits

    start = max(0, min(start, nbits))
    stop = max(0, min(stop, nbits))

    if start >= stop:
        return 0

    width = stop - start
    mask = (1 << width) - 1
    return (x >> start) & mask


def ref_slice_general(x: int, sl: slice) -> int:
    """
    Oráculo de referencia para slicing general (slow-path) sobre un entero Python.

    Aquí forzamos semántica Python completa de slicing:
    - Construimos la lista de bits del entero (LSB -> MSB) como una secuencia finita.
    - Aplicamos el slice real de Python `bits[sl]` (incluye step > 1 y step < 0).
    - Recompactamos el resultado en un nuevo entero:
      el primer bit extraído pasa a ser el bit 0 del entero resultante.

    Notas:
    - x debe ser natural (x >= 0).
    - Si sl.step == 0, lanza ValueError (igual que Python).
    - Como la secuencia de bits es finita (hasta bit_length), cualquier acceso fuera de rango
      se comporta como slicing de listas (no como bits infinitos).
    """
    if x < 0:
        raise ValueError

    if sl.step == 0:
        raise ValueError

    nbits = x.bit_length()
    bits = [(x >> i) & 1 for i in range(nbits)]  # LSB -> MSB
    picked = bits[sl]

    out = 0
    for j, b in enumerate(picked):
        out |= (b & 1) << j
    return out


def rand_nat(nbits: int) -> int:
    """
    Genera un natural aleatorio con hasta `nbits` bits.

    - Si nbits <= 0, devuelve 0.
    - Si nbits > 0, usa `random.getrandbits(nbits)` para obtener un entero uniforme en
      [0, 2**nbits).

    Se usa en tests de tipo property-based (muchas combinaciones rápidas).
    """
    if nbits <= 0:
        return 0
    return random.getrandbits(nbits)

# ----------------------------
# Numba helper: slicing fast sobre core (limbs, nbits)
# ----------------------------


@njit
def _u64_bit_length(x):
    """
    Versión Numba-friendly de bit_length para un uint64.

    Numba no puede llamar de forma directa a `int.bit_length()` sobre uint64 como en Python,
    así que calculamos el número de bits significativos a mano:
    - 0 -> 0
    - si x != 0 -> contamos cuántos shifts a la derecha hacen falta para llegar a 0.
    """
    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 palabras uint64 (limbs).

    `limbs` representa la magnitud en little-endian por palabra:
    - limbs[0] contiene bits 0..63 (LSB chunk)
    - limbs[1] contiene bits 64..127, etc.

    Esta función:
    - Busca desde el final la última palabra no nula (para eliminar ceros líderes).
    - Calcula cuántos bits significativos tiene esa palabra (bit_length uint64).
    - Devuelve el `nbits` total: (índice_palabra * 64 + bit_length(palabra_top)).

    Si todas las palabras son 0 o el array está vacío, 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
def slice_bits_fast_core(limbs, nbits, start, stop):
    """
    Extrae un rango de bits [start, stop) desde una representación por limbs (core),
    usando un algoritmo de “ventana” eficiente por palabras (fast-path).

    Contrato:
    - Se asume step = 1 (el caller ya eligió fast-path).
    - `limbs` es uint64 contiguo, little-endian por palabra.
    - `nbits` es el número de bits significativos del valor original.
    - `start` y `stop` llegan ya como enteros (sin None), pero aquí se recortan a [0, nbits].
    - El resultado se recompacta: el bit `start` pasa a ser el bit 0 del resultado.
    - Si el intervalo queda vacío, devuelve (array vacío, 0).

    Implementación (idea):
    - Localiza la palabra inicial `word0 = start // 64` y el desplazamiento dentro de palabra.
    - Copia/combina palabras sucesivas para aplicar el shift y reconstruir el bloque alineado.
    - Enmascara la última palabra para truncar a `width` bits.
    - Recalcula `out_nbits` quitando ceros líderes (trim).
    """
    if start < 0:
        start = 0
    if stop < 0:
        stop = 0
    if start > nbits:
        start = nbits
    if stop > nbits:
        stop = nbits
    if start >= stop:
        return np.empty(0, dtype=np.uint64), np.int64(0)

    width = stop - start
    out_words = (width + 63) >> 6

    word0 = start >> 6
    shift = start & 63

    out = np.empty(out_words, dtype=np.uint64)

    if shift == 0:
        for i in range(out_words):
            s = word0 + i
            out[i] = limbs[s] if s < limbs.size else np.uint64(0)
    else:
        rshift = shift
        lshift = 64 - shift
        for i in range(out_words):
            s0 = word0 + i
            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[i] = (lo >> rshift) | (hi << lshift)

    rem = width & 63
    if rem != 0:
        out[-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


# ----------------------------
# 1) Dominio y construcción
# ----------------------------

def test_domain_rejects_negative():
    with pytest.raises(ValueError):
        HugeNat(-1)


@pytest.mark.parametrize("v", [0, 1, 2, 3, 7, 8, 2 ** 63, 2 ** 127 + 12345])
def test_int_and_index_roundtrip(v):
    x = HugeNat(v)
    assert int(x) == v
    assert operator.index(x) == v
    assert bool(x) == (v != 0)


def test_hash_matches_int():
    x = HugeNat(123456789)
    assert hash(x) == hash(int(x))


def test_str_matches_int():
    v = 2 ** 80 + 7
    x = HugeNat(v)
    assert str(x) == str(v)


# ----------------------------
# 2) Paridad API int (mínimos)
# ----------------------------

@pytest.mark.parametrize("v", [0, 1, 2, 3, 255, 256, 2 ** 100 + 3])
def test_bit_length_and_count_match_int(v):
    x = HugeNat(v)
    assert x.bit_length() == v.bit_length()
    assert x.bit_count() == v.bit_count()


def test_to_bytes_from_bytes_roundtrip():
    v = 2 ** 200 + 2 ** 64 + 7
    x = HugeNat(v)
    length = (v.bit_length() + 7) // 8
    b = x.to_bytes(length=length, byteorder="big", signed=False)
    y = HugeNat.from_bytes(b, byteorder="big", signed=False)
    assert int(y) == v


# ----------------------------
# 3) Aritmética/bitwise y retorno
# ----------------------------

def test_basic_ops_with_hugenat_return_hugenat():
    a = HugeNat(10)
    b = HugeNat(7)

    r1 = a + b
    r2 = a * b
    r3 = a << 5
    r4 = a | b
    r5 = a & b
    r6 = a ^ b
    r7 = a // b
    r8 = a % b
    r9 = pow(a, 3)

    for r in (r1, r2, r3, r4, r5, r6, r7, r8, r9):
        assert isinstance(r, HugeNat)

    assert int(r1) == 17
    assert int(r2) == 70
    assert int(r3) == 10 << 5
    assert int(r4) == (10 | 7)
    assert int(r5) == (10 & 7)
    assert int(r6) == (10 ^ 7)
    assert int(r7) == (10 // 7)
    assert int(r8) == (10 % 7)
    assert int(r9) == (10 ** 3)


def test_ops_with_python_int():
    x = HugeNat(123)
    assert isinstance(x + 7, HugeNat)
    assert isinstance(7 + x, HugeNat)
    assert int(x + 7) == 130
    assert int(7 + x) == 130
    assert int(x & 0b1010) == (123 & 0b1010)
    assert int(0b1010 & x) == (0b1010 & 123)


def test_bitwise_ops_and_shifts():
    a = HugeNat(0b1100)
    b = HugeNat(0b1010)

    assert isinstance(a & b, HugeNat)
    assert isinstance(a | b, HugeNat)
    assert isinstance(a ^ b, HugeNat)

    assert int(a & b) == (0b1100 & 0b1010)
    assert int(a | b) == (0b1100 | 0b1010)
    assert int(a ^ b) == (0b1100 ^ 0b1010)

    assert isinstance(a << 2, HugeNat)
    assert isinstance(a >> 1, HugeNat)
    assert int(a << 2) == (0b1100 << 2)
    assert int(a >> 1) == (0b1100 >> 1)


def test_bit_rotations():
    x = HugeNat(0b100101)
    assert int(x.rotl(2)) == 0b010110
    assert int(x.rotr(2)) == 0b011001

    # Ancho cero devuelve 0
    z = HugeNat(0)
    assert int(z.rotl(5)) == 0
    assert int(z.rotr(5)) == 0


def test_init_from_array_uint64():
    limbs = np.array([0xFFFFFFFFFFFFFFFF, 0x1], dtype=np.uint64)
    x = HugeNat(limbs)
    assert int(x) == (1 << 64) | 0xFFFFFFFFFFFFFFFF


def test_init_from_array_ints():
    limbs = [1, 2, 3]
    x = HugeNat(limbs)
    assert int(x) == (3 << 128) | (2 << 64) | 1


def test_arrays_with_trailing_zero_limb_are_equivalent():
    a = HugeNat([0xAAAAAAAAAAAAAAAA, 0xAAAAAAAAAAAAAAAA, 0])
    b = HugeNat([0xAAAAAAAAAAAAAAAA, 0xAAAAAAAAAAAAAAAA])
    c = HugeNat([0, 0xAAAAAAAAAAAAAAAA, 0xAAAAAAAAAAAAAAAA])

    assert a == b
    assert a != c


def test_float_array_fractional_rejected():
    arr = np.array([1.5, 2.0], dtype=float)
    with pytest.raises(ValueError):
        HugeNat(arr)


def test_float_array_integer_like_accepted():
    arr = np.array([1.0, 2.0, 3.0], dtype=float)
    x = HugeNat(arr)
    assert int(x) == (3 << 128) | (2 << 64) | 1


def test_float_array_large_rejected_for_precision():
    arr = np.array([2 ** 60, 2 ** 61], dtype=float)
    with pytest.raises(ValueError):
        HugeNat(arr)


def test_subtraction_cannot_go_negative():
    with pytest.raises(ValueError):
        _ = HugeNat(3) - HugeNat(5)
    with pytest.raises(ValueError):
        _ = HugeNat(3) - 5


def test_invert_is_not_defined_for_unbounded_naturals():
    with pytest.raises(ValueError):
        _ = ~HugeNat(0)
    with pytest.raises(ValueError):
        _ = ~HugeNat(123)


# ----------------------------
# 4) Indexado por bits
# ----------------------------

@pytest.mark.parametrize("v", [0, 1, 2, 3, 0xDEADBEEF, 2 ** 130 + 5])
def test_bit_indexing_matches_reference(v: int):
    x = HugeNat(v)
    n = v.bit_length()

    if n == 0:
        assert x[0] == 0
        assert x[-1] == 0
        return

    # puntos típicos
    for i in (0, 1, 2, n - 1):
        assert x[i] == ref_bit(v, i)

    assert x[-1] == ref_bit(v, -1)
    assert x[-n] == ref_bit(v, -n)

    assert x[n] == 0
    assert x[-(n + 1)] == 0


# ----------------------------
# 5) Slicing fast-path (step None/1)
# ----------------------------

@pytest.mark.parametrize("v, sl", [(0, slice(0, 1, None)),
                                   (0b101101, slice(0, 3, None)),
                                   (0b101101, slice(2, 5, None)),
                                   (0b101101, slice(None, None, None)),
                                   (0b101101, slice(-3, None, None)),
                                   (0b101101, slice(None, -1, None)),
                                   (2 ** 200 + 12345, slice(17, 140, None)),
                                   ],
                         )
def test_slice_step1_matches_reference(v, sl):
    x = HugeNat(v)
    out = x[sl]
    assert isinstance(out, HugeNat)
    assert int(out) == ref_slice_step1(v, sl.start, sl.stop)


def test_slice_empty_returns_zero_not_error():
    x = HugeNat(0b101101)
    assert int(x[3:3]) == 0
    assert int(x[4:2]) == 0


# ----------------------------
# 6) Slicing slow-path (step arbitrario)
# ----------------------------

@pytest.mark.parametrize("v, sl", [(0, slice(None, None, -1)),
                                   (0b101101, slice(None, None, 2)),
                                   (0b101101, slice(5, None, -1)),
                                   (0b101101, slice(4, 0, -2)),
                                   (2 ** 120 + 0b101101, slice(None, None, 3)),
                                   ],
                         )
def test_slice_general_step_matches_reference(v, sl):
    x = HugeNat(v)
    out = x[sl]
    assert isinstance(out, HugeNat)
    assert int(out) == ref_slice_general(v, sl)


def test_slice_step_zero_raises():
    x = HugeNat(123)
    with pytest.raises(ValueError):
        _ = x[0:10:0]


# ----------------------------
# 7) bits() view
# ----------------------------

@pytest.mark.parametrize("v", [0, 1, 2, 3, 0b101101, 2 ** 100 + 7])
def test_bits_view_default_msb_to_lsb(v):
    x = HugeNat(v)
    arr = np.asarray(x.bits(), dtype=np.uint8)

    if v == 0:
        # contrato estricto: 0 -> vista vacía
        assert arr.size == 0
        return

    # msb->lsb: primer bit siempre 1 (MSB), último es LSB
    assert arr[0] == 1
    assert arr[-1] == (v & 1)


def test_bits_view_length_padding():
    x = HugeNat(0b101)
    arr = np.asarray(x.bits(length=8), dtype=np.uint8)
    assert arr.size == 8
    # msb->lsb con padding a 8: "00000101"
    assert arr.tolist() == [0, 0, 0, 0, 0, 1, 0, 1]


def test_bits_view_order_lsb_to_msb():
    v = 0b101101
    x = HugeNat(v)
    arr = np.asarray(x.bits(order="lsb->msb"), dtype=np.uint8)
    # lsb->msb: arr[0] es LSB
    assert arr[0] == (v & 1)
    assert arr[-1] == 1  # MSB


# ----------------------------
# 8) Propiedades aleatorias: slice step1
# ----------------------------

@pytest.mark.parametrize("nbits", [0, 1, 2, 5, 64, 65, 130, 257])
def test_property_slice_step1_random(nbits):
    v = rand_nat(nbits)
    x = HugeNat(v)
    n = v.bit_length()

    for _ in range(50):
        if n == 0:
            a = random.randint(-5, 5)
            b = random.randint(-5, 5)
        else:
            a = random.randint(-n - 5, n + 5)
            b = random.randint(-n - 5, n + 5)

        sl = slice(a, b, None)
        assert int(x[sl]) == ref_slice_step1(v, sl.start, sl.stop)


# ----------------------------
# 9) Numba: el core debe ser usable en njit
# ----------------------------

def test_to_core_contract_and_roundtrip_from_core():
    v = 2 ** 200 + 2 ** 130 + 123456789
    x = HugeNat(v)

    limbs, nbits = x.to_core()

    assert isinstance(limbs, np.ndarray)
    assert limbs.dtype == np.uint64
    assert limbs.flags["C_CONTIGUOUS"]
    assert isinstance(nbits, (np.integer, int))
    assert int(nbits) == v.bit_length()

    y = HugeNat.from_core(limbs, nbits)
    assert int(y) == v


def test_numba_core_fast_slice_matches_class_slice():
    v = 2 ** 200 + 2 ** 130 + 123456789
    x = HugeNat(v)
    limbs, nbits = x.to_core()

    # Casos variados (solo step1)
    cases = [
        (0, 16),
        (8, 80),
        (63, 129),
        (17, 140),
        (140, 17),  # vacío
        (-10, None),  # None se resuelve fuera
        (None, -1),
    ]

    n = int(nbits)
    for a, b in cases:
        sl = slice(a, b, None)

        # Normalización para pasar a core: usa el mismo criterio del contrato fast-path
        start = 0 if sl.start is None else sl.start
        stop = n if sl.stop is None else sl.stop

        if start < 0:
            start += n
        if stop < 0:
            stop += n

        out_limbs, out_nbits = slice_bits_fast_core(limbs, nbits, start, stop)
        y = HugeNat.from_core(out_limbs, out_nbits)

        assert int(y) == int(x[sl])
