"""Implementación de HugeNat basada en int para CPU."""

from __future__ import annotations

from typing import Literal, Optional, Tuple, Union

import numpy as np

_INT_LIKE = Union[int, "HugeNat"]


class HugeNat:
    """Numero natural grande con soporte de indexado y slicing de bits."""

    __slots__ = ("_value", "_fixed_nbits")

    def __init__(
            self,
            value: Union[int, np.ndarray, list, tuple, HugeNat],
            bit_length: Optional[int] = None,
    ) -> None:
        """
        Crea un HugeNat a partir de int o de un array 1D de palabras.

        Si `bit_length` se especifica, fija la anchura lógica del número:
        - si `value` tiene más bits, se trunca por arriba;
        - si tiene menos, se consideran ceros a la izquierda;
        - `bit_length()` devuelve la anchura fijada.
        """
        declared_nbits = self._normalize_declared_nbits(bit_length)

        if isinstance(value, HugeNat):
            raw_value = value._value
            if declared_nbits is None:
                declared_nbits = value._fixed_nbits
        elif isinstance(value, int):
            if value < 0:
                raise ValueError("HugeNat solo acepta int no negativos")
            raw_value = value
        else:
            limbs = self._normalize_limbs(value)
            raw_value = self._limbs_to_int(limbs)

        self._value = self._apply_declared_nbits(raw_value, declared_nbits)
        self._fixed_nbits = declared_nbits

    @staticmethod
    def _normalize_declared_nbits(bit_length: Optional[int]) -> Optional[int]:
        """Valida y normaliza el bit_length declarado por el usuario."""
        if bit_length is None:
            return None
        if isinstance(bit_length, bool) or not isinstance(bit_length, int):
            raise TypeError("bit_length debe ser int")
        if bit_length < 0:
            raise ValueError("bit_length debe ser no negativo")
        return bit_length

    @staticmethod
    def _apply_declared_nbits(value: int, fixed_nbits: Optional[int]) -> int:
        """Aplica truncado si hay un bit_length fijado."""
        if fixed_nbits is None:
            return value
        if fixed_nbits == 0:
            return 0
        if value.bit_length() <= fixed_nbits:
            return value
        return value & ((1 << fixed_nbits) - 1)

    def _declared_bit_length(self) -> int:
        """Devuelve la anchura lógica: fijada por usuario o bit_length natural."""
        if self._fixed_nbits is not None:
            return self._fixed_nbits
        return self._value.bit_length()

    def __int__(self) -> int:
        """Devuelve el entero subyacente."""
        return self._value

    def __index__(self) -> int:
        """Permite usarlo en contextos de indice."""
        return self._value

    def __bool__(self) -> bool:
        """Evalua a False si es cero, True en otro caso."""
        return self._value != 0

    def __len__(self) -> int:
        """Devuelve la anchura lógica del valor."""
        return self._declared_bit_length()

    def __hash__(self) -> int:
        """Hash consistente con int."""
        return hash(self._value)

    def __str__(self) -> str:
        """Representación en cadena igual a int."""
        return str(self._value)

    def __repr__(self) -> str:
        """Representación informativa para depuracion."""
        if self._fixed_nbits is not None:
            return f"HugeNat({self._value}, bit_length={self._fixed_nbits})"
        return f"HugeNat({self._value})"

    @staticmethod
    def _coerce_other(value: _INT_LIKE) -> int:
        """Convierte HugeNat o int a int, o delega si no corresponde."""
        if isinstance(value, HugeNat):
            return value._value
        if isinstance(value, int):
            return value
        return NotImplemented  # type: ignore[return-value]

    @staticmethod
    def _coerce_natural_other(value: _INT_LIKE) -> int:
        """Convierte HugeNat o int a int no negativo, o delega si no corresponde."""
        other_val = HugeNat._coerce_other(value)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        if other_val < 0:
            raise ValueError("HugeNat solo opera con int no negativos")
        return other_val

    def bit_length(self) -> int:
        """Devuelve la anchura lógica en bits del valor."""
        return self._declared_bit_length()

    def bit_count(self) -> int:
        """Cuenta los bits a 1 del valor."""
        return self._value.bit_count()

    def to_bytes(self, length: int, byteorder: Literal["little", "big"], signed: bool = False) -> bytes:
        """Convierte a bytes igual que int.to_bytes."""
        return self._value.to_bytes(length=length, byteorder=byteorder, signed=signed)

    @classmethod
    def from_bytes(
            cls,
            data: bytes,
            byteorder: Literal["little", "big"],
            signed: bool = False,
    ) -> "HugeNat":
        """Crea un HugeNat desde bytes igual que int.from_bytes."""
        value = int.from_bytes(data, byteorder=byteorder, signed=signed)
        return cls(value)

    def __eq__(self, other: object) -> bool:
        """Comparación de igualdad con int o HugeNat."""
        if isinstance(other, HugeNat):
            return self._value == other._value
        if isinstance(other, int):
            return self._value == other
        if isinstance(other, (np.ndarray, list, tuple)):
            limbs = self._normalize_limbs(other)
            return self._value == self._limbs_to_int(limbs)
        return False

    def __lt__(self, other: _INT_LIKE) -> bool:
        """Comparación menor que con int o HugeNat."""
        other_val = self._coerce_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return self._value < other_val

    def __le__(self, other: _INT_LIKE) -> bool:
        """Comparación menor o igual con int o HugeNat."""
        other_val = self._coerce_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return self._value <= other_val

    def __gt__(self, other: _INT_LIKE) -> bool:
        """Comparación mayor que con int o HugeNat."""
        other_val = self._coerce_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return self._value > other_val

    def __ge__(self, other: _INT_LIKE) -> bool:
        """Comparación mayor o igual con int o HugeNat."""
        other_val = self._coerce_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return self._value >= other_val

    def __add__(self, other: _INT_LIKE) -> "HugeNat":
        """Suma con int o HugeNat."""
        other_val = self._coerce_natural_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return HugeNat(self._value + other_val)

    def __radd__(self, other: _INT_LIKE) -> "HugeNat":
        """Suma reflejada."""
        return self.__add__(other)

    def __mul__(self, other: _INT_LIKE) -> "HugeNat":
        """Multiplicación con int o HugeNat."""
        other_val = self._coerce_natural_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return HugeNat(self._value * other_val)

    def __rmul__(self, other: _INT_LIKE) -> "HugeNat":
        """Multiplicacion reflejada."""
        return self.__mul__(other)

    def __floordiv__(self, other: _INT_LIKE) -> "HugeNat":
        """Division entera con int o HugeNat."""
        other_val = self._coerce_natural_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return HugeNat(self._value // other_val)

    def __rfloordiv__(self, other: _INT_LIKE) -> "HugeNat":
        """Division entera reflejada."""
        other_val = self._coerce_natural_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return HugeNat(other_val // self._value)

    def __mod__(self, other: _INT_LIKE) -> "HugeNat":
        """Modulo con int o HugeNat."""
        other_val = self._coerce_natural_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return HugeNat(self._value % other_val)

    def __rmod__(self, other: _INT_LIKE) -> "HugeNat":
        """Modulo reflejado."""
        other_val = self._coerce_natural_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return HugeNat(other_val % self._value)

    def __pow__(self, exponent: _INT_LIKE, modulo: Optional[int] = None) -> "HugeNat":
        """Potencia con exponente no negativo y opcional modulo."""
        exp_val = self._coerce_natural_other(exponent)
        if exp_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        if modulo is None:
            return HugeNat(pow(self._value, exp_val))
        return HugeNat(pow(self._value, exp_val, modulo))

    def __sub__(self, other: _INT_LIKE) -> "HugeNat":
        """Resta que no permite resultados negativos."""
        other_val = self._coerce_natural_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        result = self._value - other_val
        if result < 0:
            raise ValueError("Resultado negativo no permitido")
        return HugeNat(result)

    def __rsub__(self, other: _INT_LIKE) -> "HugeNat":
        """Resta reflejada que no permite negativos."""
        other_val = self._coerce_natural_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        result = other_val - self._value
        if result < 0:
            raise ValueError("Resultado negativo no permitido")
        return HugeNat(result)

    def __and__(self, other: _INT_LIKE) -> "HugeNat":
        """AND bit a bit."""
        other_val = self._coerce_natural_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return HugeNat(self._value & other_val)

    def __rand__(self, other: _INT_LIKE) -> "HugeNat":
        """AND reflejado."""
        return self.__and__(other)

    def __or__(self, other: _INT_LIKE) -> "HugeNat":
        """OR bit a bit."""
        other_val = self._coerce_natural_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return HugeNat(self._value | other_val)

    def __ror__(self, other: _INT_LIKE) -> "HugeNat":
        """OR reflejado."""
        return self.__or__(other)

    def __xor__(self, other: _INT_LIKE) -> "HugeNat":
        """XOR bit a bit."""
        other_val = self._coerce_natural_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return HugeNat(self._value ^ other_val)

    def __rxor__(self, other: _INT_LIKE) -> "HugeNat":
        """XOR reflejado."""
        return self.__xor__(other)

    def __lshift__(self, shift: int) -> "HugeNat":
        """Desplazamiento a la izquierda."""
        if not isinstance(shift, int):
            return NotImplemented  # type: ignore[return-value]
        if shift < 0:
            raise ValueError("Desplazamiento negativo no permitido")
        return HugeNat(self._value << shift)

    def __rlshift__(self, other: _INT_LIKE) -> "HugeNat":
        """Desplazamiento reflejado."""
        other_val = self._coerce_natural_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return HugeNat(other_val << self._value)

    def __rshift__(self, shift: int) -> "HugeNat":
        """Desplazamiento a la derecha."""
        if not isinstance(shift, int):
            return NotImplemented  # type: ignore[return-value]
        if shift < 0:
            raise ValueError("Desplazamiento negativo no permitido")
        return HugeNat(self._value >> shift)

    def __rrshift__(self, other: _INT_LIKE) -> "HugeNat":
        """Desplazamiento reflejado a la derecha."""
        other_val = self._coerce_natural_other(other)
        if other_val is NotImplemented:
            return NotImplemented  # type: ignore[return-value]
        return HugeNat(other_val >> self._value)

    def __invert__(self) -> "HugeNat":
        """Negacion bit a bit no definida para naturales sin limite."""
        raise ValueError("Invert no definido para HugeNat")

    def rotl(self, shift: int) -> "HugeNat":
        """Rota a la izquierda usando el ancho lógico de bits."""
        if not isinstance(shift, int):
            raise ValueError("shift debe ser int")
        width = self.bit_length()
        if width == 0:
            return HugeNat(0)
        shift %= width
        mask = (1 << width) - 1
        v = self._value & mask
        return HugeNat(((v << shift) | (v >> (width - shift))) & mask)

    def rotr(self, shift: int) -> "HugeNat":
        """Rota a la derecha usando el ancho lógico de bits."""
        if not isinstance(shift, int):
            raise ValueError("shift debe ser int")
        width = self.bit_length()
        if width == 0:
            return HugeNat(0)
        shift %= width
        mask = (1 << width) - 1
        v = self._value & mask
        return HugeNat(((v >> shift) | (v << (width - shift))) & mask)

    def __getitem__(self, key: Union[int, slice]) -> Union[int, "HugeNat"]:
        """Acceso por indice o slicing de bits."""
        if isinstance(key, int):
            return self._get_bit(key)
        if isinstance(key, slice):
            return self._get_slice(key)
        raise TypeError("Indice no valido")

    def split(self, n: int) -> Tuple["HugeNat", ...]:
        """
        Divide el número en `n` trozos de igual cantidad de bits.

        Cada trozo se devuelve como `HugeNat` con `bit_length` fijado al ancho
        del trozo. El orden de salida es de menor a mayor peso (LSB -> MSB).
        """
        if isinstance(n, bool) or not isinstance(n, int):
            raise TypeError("n debe ser int")
        if n < 1:
            raise ValueError("n debe ser mayor o igual que 1")

        total_bits = self.bit_length()
        if total_bits % n != 0:
            raise ValueError("La cantidad de bits debe ser divisible por n")

        chunk_bits = total_bits // n
        if chunk_bits == 0:
            return tuple(HugeNat(0, bit_length=0) for _ in range(n))

        mask = (1 << chunk_bits) - 1
        return tuple(
            HugeNat((self._value >> (chunk_bits * i)) & mask, bit_length=chunk_bits)
            for i in range(n)
        )

    def _get_bit(self, i: int) -> int:
        """Obtiene un bit por indice con LSB=0, devolviendo 0 fuera de rango."""
        nbits = self.bit_length()
        if i < 0:
            i += nbits
        if i < 0 or i >= nbits:
            return 0
        return (self._value >> i) & 1

    def _get_slice(self, sl: slice) -> "HugeNat":
        """Obtiene un slice de bits con semantica Python."""
        if sl.step == 0:
            raise ValueError("step no puede ser 0")
        if sl.step is None or sl.step == 1:
            return self._slice_fast(sl.start, sl.stop)
        return self._slice_slow(sl)

    def _slice_fast(self, start: Optional[int], stop: Optional[int]) -> "HugeNat":
        """Slice rapido para step None o 1, con relleno implícito de ceros superiores."""
        nbits = self.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, start)
        stop = max(0, stop)
        if start >= stop:
            return HugeNat(0)
        if start >= nbits:
            return HugeNat(0)

        width = stop - start
        shifted = self._value >> start
        if width >= shifted.bit_length():
            return HugeNat(shifted)

        mask = (1 << width) - 1
        return HugeNat(shifted & mask)

    def _slice_slow(self, sl: slice) -> "HugeNat":
        """Slice general con step arbitrario y repack LSB-first."""
        nbits = self.bit_length()
        step = sl.step if sl.step is not None else 1

        if step > 0:
            start, stop = self._normalize_slice_bounds_positive(nbits, sl.start, sl.stop)
        else:
            start, stop = self._normalize_slice_bounds_negative(nbits, sl.start, sl.stop)

        out = 0
        for j, i in enumerate(range(start, stop, step)):
            if i < nbits:
                out |= ((self._value >> i) & 1) << j
        return HugeNat(out)

    @staticmethod
    def _normalize_slice_bounds_positive(
            nbits: int,
            start: Optional[int],
            stop: Optional[int],
    ) -> Tuple[int, int]:
        """Normaliza límites de slice para `step > 0`."""
        start_i = 0 if start is None else start
        stop_i = nbits if stop is None else stop

        if start_i < 0:
            start_i += nbits
        if stop_i < 0:
            stop_i += nbits

        return max(0, start_i), max(0, stop_i)

    @staticmethod
    def _normalize_slice_bounds_negative(
            nbits: int,
            start: Optional[int],
            stop: Optional[int],
    ) -> Tuple[int, int]:
        """Normaliza límites de slice para `step < 0`."""
        if start is None:
            start_i = nbits - 1
        else:
            start_i = start
            if start_i < 0:
                start_i += nbits
            if start_i < 0:
                start_i = -1

        if stop is None:
            stop_i = -1
        else:
            stop_i = stop
            if stop_i < 0:
                stop_i += nbits
            if stop_i < 0:
                stop_i = -1

        return start_i, stop_i

    def bits(self, order: str = "msb->lsb", length: Optional[int] = None) -> np.ndarray:
        """Devuelve una vista 1D de bits como uint8."""
        if length is not None and length < 0:
            raise ValueError("length debe ser no negativo")
        nbits = self.bit_length()
        if nbits == 0 and length is None:
            return np.empty(0, dtype=np.uint8)

        if order not in ("msb->lsb", "lsb->msb"):
            raise ValueError("order no valido")

        if nbits == 0:
            bits_list = []
        elif order == "msb->lsb":
            bits_list = [((self._value >> i) & 1) for i in range(nbits - 1, -1, -1)]
        else:
            bits_list = [((self._value >> i) & 1) for i in range(nbits)]

        if length is not None:
            if length > len(bits_list):
                pad = [0] * (length - len(bits_list))
                if order == "msb->lsb":
                    bits_list = pad + bits_list
                else:
                    bits_list += pad
            elif length < len(bits_list):
                if order == "msb->lsb":
                    bits_list = bits_list[-length:]
                else:
                    bits_list = bits_list[:length]

        return np.asarray(bits_list, dtype=np.uint8)

    def bits_str(
            self,
            order: str = "msb->lsb",
            group: int = 64,
            sep: str = " ",
    ) -> str:
        """Devuelve una representacion en string con bits agrupados."""
        if group <= 0:
            raise ValueError("group debe ser mayor que 0")
        if order not in ("msb->lsb", "lsb->msb"):
            raise ValueError("order no valido")

        nbits = self.bit_length()
        if nbits == 0:
            return ""

        if order == "msb->lsb":
            bits_list = [((self._value >> i) & 1) for i in range(nbits - 1, -1, -1)]
        else:
            bits_list = [((self._value >> i) & 1) for i in range(nbits)]

        chars = [str(b) for b in bits_list]
        groups = ["".join(chars[i:i + group]) for i in range(0, len(chars), group)]
        return sep.join(groups)

    def to_core(self) -> Tuple[np.ndarray, np.int64]:
        """Devuelve (limbs, nbits) en formato compatible con Numba."""
        nbits = self.bit_length()
        if nbits == 0:
            limbs = np.empty(0, dtype=np.uint64)
            return limbs, np.int64(0)

        n_words = (nbits + 63) >> 6
        limbs = np.empty(n_words, dtype=np.uint64)
        mask = (1 << 64) - 1
        for i in range(n_words):
            limbs[i] = np.uint64((self._value >> (64 * i)) & mask)
        return limbs, np.int64(nbits)

    @classmethod
    def from_core(cls, limbs: np.ndarray, nbits: Union[np.integer, int]) -> "HugeNat":
        """Construye un HugeNat desde (limbs, nbits) con anchura lógica explícita."""
        if not isinstance(limbs, np.ndarray):
            raise TypeError("limbs debe ser np.ndarray")
        if limbs.ndim != 1:
            raise ValueError("limbs debe ser 1D")
        if limbs.dtype != np.uint64:
            raise ValueError("limbs debe tener dtype np.uint64")
        if not limbs.flags["C_CONTIGUOUS"]:
            raise ValueError("limbs debe ser contiguo en memoria")
        if isinstance(nbits, bool) or not isinstance(nbits, (int, np.integer)):
            raise TypeError("nbits debe ser int")

        nbits_int = int(nbits)
        if nbits_int < 0:
            raise ValueError("nbits debe ser no negativo")
        if nbits_int == 0:
            if limbs.size != 0:
                raise ValueError("limbs debe ser vacio cuando nbits es 0")
            return cls(0, bit_length=0)

        expected_words = (nbits_int + 63) >> 6
        if limbs.size != expected_words:
            raise ValueError("limbs y nbits no son consistentes")

        last_bits = (nbits_int - 1) & 63
        last_mask = (1 << (last_bits + 1)) - 1
        if (int(limbs[-1]) & ~last_mask) != 0:
            raise ValueError("El limb mas significativo contiene bits fuera de nbits")

        value = 0
        for i, limb in enumerate(limbs):
            value |= int(limb) << (64 * i)

        return cls(value, bit_length=nbits_int)

    @staticmethod
    def _normalize_limbs(value) -> np.ndarray:
        """
        Normaliza limbs a np.uint64 1D little-endian y recorta ceros líderes (MS).
        Acepta floats solo si son enteros exactos y no negativos.
        """
        if isinstance(value, (list, tuple)):
            arr = np.asarray(value, dtype=object)
        else:
            arr = np.asarray(value)

        if arr.ndim != 1:
            raise ValueError("El array de limbs debe ser 1D")

        if arr.size == 0:
            return np.empty(0, dtype=np.uint64)

        # Floats: aceptar solo si son enteros exactos, no negativos y precisos (<= 2**53)
        if arr.dtype.kind == "f":
            if not np.all(np.isfinite(arr)):
                raise ValueError("El array debe contener valores finitos")
            if not np.all(arr == np.floor(arr)):
                raise ValueError("El array debe contener enteros")
            if np.any(arr < 0):
                raise ValueError("El array debe contener enteros no negativos")
            if np.any(arr > 2**53):
                raise ValueError("Floats demasiado grandes para representacion exacta")
            if np.any(arr > 0xFFFFFFFFFFFFFFFF):
                raise ValueError("Cada limb debe caber en 64 bits")
            arr_u64 = arr.astype(np.uint64, copy=False)
            return HugeNat._trim_high_zero_limbs(arr_u64)

        # Object: validar elemento a elemento (exactitud)
        if arr.dtype == object:
            out = np.empty(arr.size, dtype=np.uint64)
            for i, v in enumerate(arr):
                if not isinstance(v, (int, np.integer, float, np.floating)):
                    raise ValueError("El array debe contener enteros")
                if isinstance(v, (float, np.floating)) and v != int(v):
                    raise ValueError("El array debe contener enteros")
                if isinstance(v, (float, np.floating)) and v > 2**53:
                    raise ValueError("Floats demasiado grandes para representacion exacta")
                iv = int(v)
                if iv < 0:
                    raise ValueError("El array debe contener enteros no negativos")
                if iv > 0xFFFFFFFFFFFFFFFF:
                    raise ValueError("Cada limb debe caber en 64 bits")
                out[i] = np.uint64(iv)
            arr_u64 = out
        else:
            # Enteros (signed/unsigned)
            if arr.dtype.kind not in ("i", "u"):
                raise ValueError("El array debe contener enteros")
            if arr.dtype.kind == "i" and np.any(arr < 0):
                raise ValueError("El array debe contener enteros no negativos")
            # Cast seguro (no oculta signo; ya validado)
            arr_u64 = arr.astype(np.uint64, copy=False)

        return HugeNat._trim_high_zero_limbs(arr_u64)

    @staticmethod
    def _trim_high_zero_limbs(arr_u64: np.ndarray) -> np.ndarray:
        """Recorta limbs de mayor peso en cero y asegura contigüidad."""
        j = arr_u64.size
        while j > 0 and arr_u64[j - 1] == 0:
            j -= 1
        if j == 0:
            return np.empty(0, dtype=np.uint64)
        return np.ascontiguousarray(arr_u64[:j])

    @staticmethod
    def _limbs_to_int(limbs: np.ndarray) -> int:
        """Reconstruye el entero a partir de limbs little-endian."""
        value = 0
        for i, limb in enumerate(limbs):
            value |= int(limb) << (64 * i)
        return value


__all__ = ["HugeNat"]
