"""
helix_cdc.native.pistons - Python ctypes bindings for libhelix_pistons.so

Usage:
    from helix_cdc.native import entropy_bytes, sponge_full, is_available

    if is_available():
        entropy = entropy_bytes(data)
    else:
        # Fallback to Python implementation
        entropy = python_entropy(data)
"""

import ctypes
import ctypes.util
import numpy as np
from pathlib import Path
from typing import Optional, Union

# Library paths to search
LIB_PATHS = [
    Path(__file__).parent / "libhelix_pistons.so",
    Path(__file__).parent.parent.parent / "c" / "pistons" / "libhelix_pistons.so",
    Path("/usr/local/lib/libhelix_pistons.so"),
]

_lib: Optional[ctypes.CDLL] = None
_available = False


def _load_library() -> Optional[ctypes.CDLL]:
    """Load the pistons library."""
    global _lib, _available

    if _lib is not None:
        return _lib

    for path in LIB_PATHS:
        if path.exists():
            try:
                _lib = ctypes.CDLL(str(path))
                _setup_functions()
                _available = True
                return _lib
            except OSError:
                continue

    # Try system library path
    lib_path = ctypes.util.find_library("helix_pistons")
    if lib_path:
        try:
            _lib = ctypes.CDLL(lib_path)
            _setup_functions()
            _available = True
            return _lib
        except OSError:
            pass

    return None


def _setup_functions():
    """Set up function signatures."""
    global _lib

    # helix_entropy_bytes
    _lib.helix_entropy_bytes.argtypes = [ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t]
    _lib.helix_entropy_bytes.restype = ctypes.c_double

    # helix_entropy_batch
    _lib.helix_entropy_batch.argtypes = [
        ctypes.POINTER(ctypes.c_uint8),
        ctypes.c_size_t,
        ctypes.c_size_t,
        ctypes.POINTER(ctypes.c_double),
    ]
    _lib.helix_entropy_batch.restype = ctypes.c_size_t

    # helix_sponge_absorb
    _lib.helix_sponge_absorb.argtypes = [
        ctypes.POINTER(ctypes.c_uint8),
        ctypes.POINTER(ctypes.c_uint8),
        ctypes.c_size_t,
    ]
    _lib.helix_sponge_absorb.restype = None

    # helix_sponge_permute
    _lib.helix_sponge_permute.argtypes = [ctypes.POINTER(ctypes.c_uint8)]
    _lib.helix_sponge_permute.restype = None

    # helix_sponge_full
    _lib.helix_sponge_full.argtypes = [
        ctypes.POINTER(ctypes.c_uint8),
        ctypes.POINTER(ctypes.c_uint8),
        ctypes.c_size_t,
        ctypes.c_size_t,
    ]
    _lib.helix_sponge_full.restype = None

    # helix_regen_step
    _lib.helix_regen_step.argtypes = [
        ctypes.POINTER(ctypes.c_uint8),
        ctypes.c_size_t,
        ctypes.POINTER(ctypes.c_uint8),
        ctypes.c_size_t,
        ctypes.c_uint32,
    ]
    _lib.helix_regen_step.restype = ctypes.c_int

    # helix_sha256_chain_step (matches Python substrate API exactly)
    _lib.helix_sha256_chain_step.argtypes = [
        ctypes.POINTER(ctypes.c_uint8),  # state_in (32 bytes)
        ctypes.POINTER(ctypes.c_uint8),  # state_out (32 bytes)
        ctypes.c_uint32,                 # iteration
    ]
    _lib.helix_sha256_chain_step.restype = ctypes.c_int

    # helix_pistons_version
    _lib.helix_pistons_version.argtypes = []
    _lib.helix_pistons_version.restype = ctypes.c_char_p


def is_available() -> bool:
    """Check if native pistons library is available."""
    _load_library()
    return _available


def version() -> str:
    """Get library version."""
    lib = _load_library()
    if lib is None:
        return "not available"
    return lib.helix_pistons_version().decode("utf-8")


def entropy_bytes(data: Union[bytes, np.ndarray]) -> float:
    """
    Calculate Shannon entropy of byte data.

    Args:
        data: bytes or uint8 numpy array

    Returns:
        Entropy in bits (0.0 to 8.0)
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("libhelix_pistons.so not available")

    if isinstance(data, np.ndarray):
        data = data.astype(np.uint8)
        data_ptr = data.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
        length = len(data)
    else:
        data_arr = (ctypes.c_uint8 * len(data)).from_buffer_copy(data)
        data_ptr = ctypes.cast(data_arr, ctypes.POINTER(ctypes.c_uint8))
        length = len(data)

    return lib.helix_entropy_bytes(data_ptr, length)


def entropy_batch(data: np.ndarray, sample_size: int) -> np.ndarray:
    """
    Calculate entropy for multiple samples.

    Args:
        data: Flat uint8 array containing all samples
        sample_size: Size of each sample

    Returns:
        Array of entropy values
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("libhelix_pistons.so not available")

    data = data.astype(np.uint8).flatten()
    n_samples = len(data) // sample_size

    data_ptr = data.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
    result = np.zeros(n_samples, dtype=np.float64)
    result_ptr = result.ctypes.data_as(ctypes.POINTER(ctypes.c_double))

    lib.helix_entropy_batch(data_ptr, sample_size, n_samples, result_ptr)

    return result


def sponge_absorb(state: np.ndarray, message: bytes) -> None:
    """
    Sponge absorb: XOR message into state and permute.

    Args:
        state: uint8 array of size 200 (modified in place)
        message: bytes to absorb
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("libhelix_pistons.so not available")

    if len(state) != 200:
        raise ValueError("State must be 200 bytes")

    state_ptr = state.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
    msg_arr = (ctypes.c_uint8 * len(message)).from_buffer_copy(message)
    msg_ptr = ctypes.cast(msg_arr, ctypes.POINTER(ctypes.c_uint8))

    lib.helix_sponge_absorb(state_ptr, msg_ptr, len(message))


def sponge_permute(state: np.ndarray) -> None:
    """
    Sponge permutation (in-place).

    Args:
        state: uint8 array of size 200 (modified in place)
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("libhelix_pistons.so not available")

    if len(state) != 200:
        raise ValueError("State must be 200 bytes")

    state_ptr = state.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
    lib.helix_sponge_permute(state_ptr)


def sponge_full(message: bytes, n_rounds: int = 100) -> bytes:
    """
    Full sponge operation: absorb message and run n_rounds.

    Args:
        message: bytes to absorb
        n_rounds: number of permutation rounds

    Returns:
        Final 200-byte state
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("libhelix_pistons.so not available")

    state = np.zeros(200, dtype=np.uint8)
    state_ptr = state.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))

    msg_arr = (ctypes.c_uint8 * len(message)).from_buffer_copy(message)
    msg_ptr = ctypes.cast(msg_arr, ctypes.POINTER(ctypes.c_uint8))

    lib.helix_sponge_full(state_ptr, msg_ptr, len(message), n_rounds)

    return bytes(state)


def regen_step(seed: bytes, chunk_size: int, iteration: int = 0) -> bytes:
    """
    SHA256-based regen step: expand seed into chunk.

    Args:
        seed: seed bytes
        chunk_size: size of output chunk
        iteration: iteration number

    Returns:
        Generated chunk
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("libhelix_pistons.so not available")

    seed_arr = (ctypes.c_uint8 * len(seed)).from_buffer_copy(seed)
    seed_ptr = ctypes.cast(seed_arr, ctypes.POINTER(ctypes.c_uint8))

    chunk = (ctypes.c_uint8 * chunk_size)()
    chunk_ptr = ctypes.cast(chunk, ctypes.POINTER(ctypes.c_uint8))

    result = lib.helix_regen_step(seed_ptr, len(seed), chunk_ptr, chunk_size, iteration)

    if result != 0:
        raise RuntimeError("regen_step failed")

    return bytes(chunk)


def sha256_chain_step(state_in: bytes, iteration: int) -> bytes:
    """
    Single SHA256 hash chain step (matches Python exactly).
    state_out = SHA256(state_in || iteration_le)

    This is the per-iteration step used in substrate API regen().

    Args:
        state_in: 32-byte input state
        iteration: iteration number

    Returns:
        32-byte output state
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("libhelix_pistons.so not available")

    if len(state_in) != 32:
        raise ValueError("state_in must be 32 bytes")

    state_in_arr = (ctypes.c_uint8 * 32).from_buffer_copy(state_in)
    state_in_ptr = ctypes.cast(state_in_arr, ctypes.POINTER(ctypes.c_uint8))

    state_out = (ctypes.c_uint8 * 32)()
    state_out_ptr = ctypes.cast(state_out, ctypes.POINTER(ctypes.c_uint8))

    result = lib.helix_sha256_chain_step(state_in_ptr, state_out_ptr, iteration)

    if result != 0:
        raise RuntimeError("sha256_chain_step failed")

    return bytes(state_out)


# Try to load on import
_load_library()
