"""
CDNA Fused Decode Kernel - Python ctypes Wrapper
WO-CDNA-CPP-KERNEL-01 + WO-CDNA2-ENCODER-KERNEL-01 + WO-CDNA2-KERNEL-02

Provides Python interface to libcdna_kernel.so for:
- SIMD-accelerated codebook gather (AVX2)
- SIMD-accelerated centroid assignment (AVX2) [encoder]
- Brotli decompression
- Zstd decompression (WO-CDNA2-KERNEL-02 Phase 1: ~29% faster than brotli)
- Fused decode pipeline (brotli and zstd variants)
- Fused decode+matmul (WO-CDNA2-KERNEL-02 Phase 2: eliminates Python/C++ boundary)

Usage:
    from helix_cdc.native.cdna_kernel import (
        cpp_kernel_available,
        fused_decode_block,         # brotli version
        fused_decode_block_zstd,    # zstd version (WO-CDNA2-KERNEL-02)
        fused_decode_matmul,        # Phase 2: decode + BLAS matmul
        gather_codebook,
        assign_nearest,  # encoder function
    )

    if cpp_kernel_available():
        # Brotli-compressed data
        W_block = fused_decode_block(compressed, codebook, rows, cols)
        # Zstd-compressed data (faster)
        W_block = fused_decode_block_zstd(compressed, codebook, rows, cols)
        # Fused decode+matmul (Phase 2) - Y accumulator modified in-place
        fused_decode_matmul(compressed, codebook, X, Y, block_rows, cols)

    # Encoder path (opt-in via env var)
    if cpp_encoder_available():
        indices = assign_nearest(values, codebook)
"""

import ctypes
import os
from ctypes import c_int, c_size_t, c_uint8, c_uint32, c_uint64, c_float, c_char_p
from ctypes import POINTER, byref
from pathlib import Path
from typing import Optional, Tuple
import numpy as np

# WO-CDNA2-ENCODER-KERNEL-01: Opt-in C++ encoder via environment variable
_USE_CPP_ENCODER = os.getenv("HELIX_USE_CPP_ENCODER", "0") == "1"

# Library paths to search (in order)
_LIB_SEARCH_PATHS = [
    Path(__file__).parent / "libcdna_kernel.so",  # Installed location
    Path(__file__).parent.parent.parent / "cpp" / "cdna_kernel" / "libcdna_kernel.so",  # Build location
]

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


def _find_library() -> Optional[Path]:
    """Find libcdna_kernel.so in known locations."""
    for path in _LIB_SEARCH_PATHS:
        if path.exists():
            return path
    return None


def _load_library() -> Optional[ctypes.CDLL]:
    """Load the C++ kernel library."""
    global _lib, _lib_load_attempted

    if _lib_load_attempted:
        return _lib

    _lib_load_attempted = True

    lib_path = _find_library()
    if lib_path is None:
        return None

    try:
        _lib = ctypes.CDLL(str(lib_path))
        _setup_argtypes(_lib)
        return _lib
    except OSError:
        return None


def _setup_argtypes(lib: ctypes.CDLL) -> None:
    """Configure function signatures."""

    # const char* cdna_kernel_version(void)
    lib.cdna_kernel_version.argtypes = []
    lib.cdna_kernel_version.restype = c_char_p

    # int cdna_has_avx2(void)
    lib.cdna_has_avx2.argtypes = []
    lib.cdna_has_avx2.restype = c_int

    # int cdna_has_avx512(void)
    lib.cdna_has_avx512.argtypes = []
    lib.cdna_has_avx512.restype = c_int

    # int cdna_brotli_decompress(...)
    lib.cdna_brotli_decompress.argtypes = [
        POINTER(c_uint8),  # compressed_data
        c_size_t,          # compressed_len
        POINTER(c_uint8),  # output
        c_size_t,          # output_capacity
        POINTER(c_size_t), # output_len
    ]
    lib.cdna_brotli_decompress.restype = c_int

    # int cdna_zstd_decompress(...) - WO-CDNA2-KERNEL-02
    lib.cdna_zstd_decompress.argtypes = [
        POINTER(c_uint8),  # compressed_data
        c_size_t,          # compressed_len
        POINTER(c_uint8),  # output
        c_size_t,          # output_capacity
        POINTER(c_size_t), # output_len
    ]
    lib.cdna_zstd_decompress.restype = c_int

    # int cdna_gather_codebook(...)
    lib.cdna_gather_codebook.argtypes = [
        POINTER(c_uint8),  # indices
        POINTER(c_float),  # codebook
        POINTER(c_float),  # output
        c_size_t,          # n_elements
    ]
    lib.cdna_gather_codebook.restype = c_int

    # int cdna_gather_codebook_2d(...)
    lib.cdna_gather_codebook_2d.argtypes = [
        POINTER(c_uint8),  # indices
        POINTER(c_float),  # codebook
        POINTER(c_float),  # output
        c_size_t,          # rows
        c_size_t,          # cols
    ]
    lib.cdna_gather_codebook_2d.restype = c_int

    # int cdna_apply_sidecar(...)
    lib.cdna_apply_sidecar.argtypes = [
        POINTER(c_float),   # weights
        c_size_t,           # rows
        c_size_t,           # cols
        POINTER(c_uint32),  # sidecar_rows
        POINTER(c_uint32),  # sidecar_cols
        POINTER(c_float),   # sidecar_values
        c_size_t,           # n_sidecar
    ]
    lib.cdna_apply_sidecar.restype = c_int

    # int cdna_apply_sidecar_linear(...)
    lib.cdna_apply_sidecar_linear.argtypes = [
        POINTER(c_float),   # weights
        c_size_t,           # n_elements
        POINTER(c_uint64),  # sidecar_positions
        POINTER(c_float),   # sidecar_values
        c_size_t,           # n_sidecar
    ]
    lib.cdna_apply_sidecar_linear.restype = c_int

    # int cdna_fused_decode_block(...)
    lib.cdna_fused_decode_block.argtypes = [
        POINTER(c_uint8),   # compressed_data
        c_size_t,           # compressed_len
        POINTER(c_float),   # codebook
        POINTER(c_float),   # output_weights
        c_size_t,           # rows
        c_size_t,           # cols
        POINTER(c_uint64),  # sidecar_positions (nullable)
        POINTER(c_float),   # sidecar_values (nullable)
        c_size_t,           # n_sidecar
        POINTER(c_uint8),   # verify_sha256 (nullable)
    ]
    lib.cdna_fused_decode_block.restype = c_int

    # int cdna_fused_decode_block_zstd(...) - WO-CDNA2-KERNEL-02
    lib.cdna_fused_decode_block_zstd.argtypes = [
        POINTER(c_uint8),   # compressed_data
        c_size_t,           # compressed_len
        POINTER(c_float),   # codebook
        POINTER(c_float),   # output_weights
        c_size_t,           # rows
        c_size_t,           # cols
        POINTER(c_uint64),  # sidecar_positions (nullable)
        POINTER(c_float),   # sidecar_values (nullable)
        c_size_t,           # n_sidecar
        POINTER(c_uint8),   # verify_sha256 (nullable)
    ]
    lib.cdna_fused_decode_block_zstd.restype = c_int

    # int cdna_sha256(...)
    lib.cdna_sha256.argtypes = [
        POINTER(c_uint8),  # data
        c_size_t,          # len
        POINTER(c_uint8),  # hash_out
    ]
    lib.cdna_sha256.restype = c_int

    # int cdna_sha256_equal(...)
    lib.cdna_sha256_equal.argtypes = [
        POINTER(c_uint8),  # hash1
        POINTER(c_uint8),  # hash2
    ]
    lib.cdna_sha256_equal.restype = c_int

    # WO-CDNA2-KERNEL-02 Phase 2: Fused decode + matmul functions

    # int cdna_fused_decode_matmul(...) - brotli version
    lib.cdna_fused_decode_matmul.argtypes = [
        POINTER(c_uint8),   # compressed_data
        c_size_t,           # compressed_len
        POINTER(c_float),   # codebook
        POINTER(c_float),   # X
        POINTER(c_float),   # Y_partial
        c_size_t,           # batch_seq
        c_size_t,           # block_rows
        c_size_t,           # cols
        POINTER(c_uint64),  # sidecar_positions (nullable)
        POINTER(c_float),   # sidecar_values (nullable)
        c_size_t,           # n_sidecar
        c_int,              # accumulate
    ]
    lib.cdna_fused_decode_matmul.restype = c_int

    # int cdna_fused_decode_matmul_zstd(...) - zstd version
    lib.cdna_fused_decode_matmul_zstd.argtypes = [
        POINTER(c_uint8),   # compressed_data
        c_size_t,           # compressed_len
        POINTER(c_float),   # codebook
        POINTER(c_float),   # X
        POINTER(c_float),   # Y_partial
        c_size_t,           # batch_seq
        c_size_t,           # block_rows
        c_size_t,           # cols
        POINTER(c_uint64),  # sidecar_positions (nullable)
        POINTER(c_float),   # sidecar_values (nullable)
        c_size_t,           # n_sidecar
        c_int,              # accumulate
    ]
    lib.cdna_fused_decode_matmul_zstd.restype = c_int

    # int cdna_fused_decode_matmul_transpose(...) - transpose version
    lib.cdna_fused_decode_matmul_transpose.argtypes = [
        POINTER(c_uint8),   # compressed_data
        c_size_t,           # compressed_len
        POINTER(c_float),   # codebook
        POINTER(c_float),   # X
        POINTER(c_float),   # Y_partial
        c_size_t,           # batch_seq
        c_size_t,           # block_rows
        c_size_t,           # cols
        POINTER(c_uint64),  # sidecar_positions (nullable)
        POINTER(c_float),   # sidecar_values (nullable)
        c_size_t,           # n_sidecar
        c_int,              # accumulate
        c_int,              # use_zstd
    ]
    lib.cdna_fused_decode_matmul_transpose.restype = c_int

    # WO-CDNA2-ENCODER-KERNEL-01: Encoder functions

    # const char* cdna_encoder_version(void)
    lib.cdna_encoder_version.argtypes = []
    lib.cdna_encoder_version.restype = c_char_p

    # int cdna_assign_nearest(...)
    lib.cdna_assign_nearest.argtypes = [
        POINTER(c_float),  # values
        POINTER(c_float),  # codebook
        POINTER(c_uint8),  # indices_out
        c_size_t,          # n_elements
    ]
    lib.cdna_assign_nearest.restype = c_int


def cpp_kernel_available() -> bool:
    """Check if the C++ kernel is available."""
    return _load_library() is not None


def get_version() -> Optional[str]:
    """Get the kernel version string."""
    lib = _load_library()
    if lib is None:
        return None
    version_bytes = lib.cdna_kernel_version()
    return version_bytes.decode('utf-8') if version_bytes else None


def has_avx2() -> bool:
    """Check if AVX2 is supported."""
    lib = _load_library()
    if lib is None:
        return False
    return lib.cdna_has_avx2() == 1


def has_avx512() -> bool:
    """Check if AVX-512 is supported."""
    lib = _load_library()
    if lib is None:
        return False
    return lib.cdna_has_avx512() == 1


def brotli_decompress(compressed: bytes, output_capacity: int) -> bytes:
    """
    Decompress brotli data using C++ kernel.

    Args:
        compressed: Brotli-compressed bytes
        output_capacity: Expected decompressed size

    Returns:
        Decompressed bytes

    Raises:
        RuntimeError: If decompression fails or kernel unavailable
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("C++ kernel not available")

    output = (c_uint8 * output_capacity)()
    output_len = c_size_t(0)

    compressed_arr = (c_uint8 * len(compressed)).from_buffer_copy(compressed)

    ret = lib.cdna_brotli_decompress(
        compressed_arr,
        len(compressed),
        output,
        output_capacity,
        byref(output_len),
    )

    if ret != 0:
        raise RuntimeError(f"Brotli decompression failed (ret={ret})")

    return bytes(output[:output_len.value])


def gather_codebook(
    indices: np.ndarray,
    codebook: np.ndarray,
) -> np.ndarray:
    """
    Gather values from codebook using uint8 indices.
    Uses AVX2 SIMD acceleration when available.

    Args:
        indices: uint8 array of codebook indices
        codebook: float32 array of shape (256,)

    Returns:
        float32 array of gathered values (same shape as indices)
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("C++ kernel not available")

    # Ensure correct dtypes and contiguity
    indices = np.ascontiguousarray(indices, dtype=np.uint8)
    codebook = np.ascontiguousarray(codebook, dtype=np.float32)

    if codebook.shape != (256,):
        raise ValueError(f"Codebook must have shape (256,), got {codebook.shape}")

    # Allocate output
    output = np.empty(indices.shape, dtype=np.float32)
    n_elements = indices.size

    # Get ctypes pointers
    indices_ptr = indices.ctypes.data_as(POINTER(c_uint8))
    codebook_ptr = codebook.ctypes.data_as(POINTER(c_float))
    output_ptr = output.ctypes.data_as(POINTER(c_float))

    ret = lib.cdna_gather_codebook(
        indices_ptr,
        codebook_ptr,
        output_ptr,
        n_elements,
    )

    if ret != 0:
        raise RuntimeError(f"Codebook gather failed (ret={ret})")

    return output


def apply_sidecar_linear(
    weights: np.ndarray,
    positions: np.ndarray,
    values: np.ndarray,
) -> None:
    """
    Apply sidecar corrections in-place using linear positions.

    Args:
        weights: float32 array to modify (will be modified in-place)
        positions: uint64 array of linear positions
        values: float32 array of correction values
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("C++ kernel not available")

    weights = np.ascontiguousarray(weights, dtype=np.float32)
    positions = np.ascontiguousarray(positions, dtype=np.uint64)
    values = np.ascontiguousarray(values, dtype=np.float32)

    weights_ptr = weights.ctypes.data_as(POINTER(c_float))
    positions_ptr = positions.ctypes.data_as(POINTER(c_uint64))
    values_ptr = values.ctypes.data_as(POINTER(c_float))

    ret = lib.cdna_apply_sidecar_linear(
        weights_ptr,
        weights.size,
        positions_ptr,
        values_ptr,
        len(positions),
    )

    if ret != 0:
        raise RuntimeError(f"Sidecar apply failed (ret={ret})")


def fused_decode_block(
    compressed_data: bytes,
    codebook: np.ndarray,
    rows: int,
    cols: int,
    sidecar_positions: Optional[np.ndarray] = None,
    sidecar_values: Optional[np.ndarray] = None,
    verify_sha256: Optional[bytes] = None,
) -> np.ndarray:
    """
    Fused decode: decompress + gather + sidecar in one call.

    This is the main hot-path function that replaces Python's:
    - brotli.decompress()
    - codebook[indices]
    - sidecar application

    Args:
        compressed_data: Brotli-compressed CDNA indices
        codebook: float32 array of shape (256,)
        rows: Number of rows in output
        cols: Number of columns in output
        sidecar_positions: Linear positions for corrections (optional)
        sidecar_values: Correction values (optional)
        verify_sha256: Expected SHA256 of decompressed indices (optional)

    Returns:
        float32 array of shape (rows, cols)

    Raises:
        RuntimeError: On decompression error or SHA256 mismatch
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("C++ kernel not available")

    codebook = np.ascontiguousarray(codebook, dtype=np.float32)
    if codebook.shape != (256,):
        raise ValueError(f"Codebook must have shape (256,), got {codebook.shape}")

    # Allocate output
    output = np.empty((rows, cols), dtype=np.float32)

    # Prepare compressed data
    compressed_arr = (c_uint8 * len(compressed_data)).from_buffer_copy(compressed_data)

    # Prepare sidecar (nullable)
    if sidecar_positions is not None and sidecar_values is not None:
        sidecar_positions = np.ascontiguousarray(sidecar_positions, dtype=np.uint64)
        sidecar_values = np.ascontiguousarray(sidecar_values, dtype=np.float32)
        positions_ptr = sidecar_positions.ctypes.data_as(POINTER(c_uint64))
        values_ptr = sidecar_values.ctypes.data_as(POINTER(c_float))
        n_sidecar = len(sidecar_positions)
    else:
        positions_ptr = None
        values_ptr = None
        n_sidecar = 0

    # Prepare SHA256 (nullable)
    if verify_sha256 is not None:
        sha256_arr = (c_uint8 * 32).from_buffer_copy(verify_sha256)
        sha256_ptr = sha256_arr
    else:
        sha256_ptr = None

    # Get pointers
    codebook_ptr = codebook.ctypes.data_as(POINTER(c_float))
    output_ptr = output.ctypes.data_as(POINTER(c_float))

    ret = lib.cdna_fused_decode_block(
        compressed_arr,
        len(compressed_data),
        codebook_ptr,
        output_ptr,
        rows,
        cols,
        positions_ptr,
        values_ptr,
        n_sidecar,
        sha256_ptr,
    )

    if ret == -1:
        raise RuntimeError("Fused decode failed (decompression or internal error)")
    elif ret == -2:
        raise RuntimeError("Fused decode failed (SHA256 mismatch)")
    elif ret != 0:
        raise RuntimeError(f"Fused decode failed (ret={ret})")

    return output


def fused_decode_block_zstd(
    compressed_data: bytes,
    codebook: np.ndarray,
    rows: int,
    cols: int,
    sidecar_positions: Optional[np.ndarray] = None,
    sidecar_values: Optional[np.ndarray] = None,
    verify_sha256: Optional[bytes] = None,
) -> np.ndarray:
    """
    Fused decode with zstd: decompress + gather + sidecar in one call.
    WO-CDNA2-KERNEL-02: ~29% faster decompress vs brotli version.

    Args:
        compressed_data: Zstd-compressed CDNA indices
        codebook: float32 array of shape (256,)
        rows: Number of rows in output
        cols: Number of columns in output
        sidecar_positions: Linear positions for corrections (optional)
        sidecar_values: Correction values (optional)
        verify_sha256: Expected SHA256 of decompressed indices (optional)

    Returns:
        float32 array of shape (rows, cols)

    Raises:
        RuntimeError: On decompression error or SHA256 mismatch
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("C++ kernel not available")

    codebook = np.ascontiguousarray(codebook, dtype=np.float32)
    if codebook.shape != (256,):
        raise ValueError(f"Codebook must have shape (256,), got {codebook.shape}")

    # Allocate output
    output = np.empty((rows, cols), dtype=np.float32)

    # Prepare compressed data
    compressed_arr = (c_uint8 * len(compressed_data)).from_buffer_copy(compressed_data)

    # Prepare sidecar (nullable)
    if sidecar_positions is not None and sidecar_values is not None:
        sidecar_positions = np.ascontiguousarray(sidecar_positions, dtype=np.uint64)
        sidecar_values = np.ascontiguousarray(sidecar_values, dtype=np.float32)
        positions_ptr = sidecar_positions.ctypes.data_as(POINTER(c_uint64))
        values_ptr = sidecar_values.ctypes.data_as(POINTER(c_float))
        n_sidecar = len(sidecar_positions)
    else:
        positions_ptr = None
        values_ptr = None
        n_sidecar = 0

    # Prepare SHA256 (nullable)
    if verify_sha256 is not None:
        sha256_arr = (c_uint8 * 32).from_buffer_copy(verify_sha256)
        sha256_ptr = sha256_arr
    else:
        sha256_ptr = None

    # Get pointers
    codebook_ptr = codebook.ctypes.data_as(POINTER(c_float))
    output_ptr = output.ctypes.data_as(POINTER(c_float))

    ret = lib.cdna_fused_decode_block_zstd(
        compressed_arr,
        len(compressed_data),
        codebook_ptr,
        output_ptr,
        rows,
        cols,
        positions_ptr,
        values_ptr,
        n_sidecar,
        sha256_ptr,
    )

    if ret == -1:
        raise RuntimeError("Fused decode (zstd) failed (decompression or internal error)")
    elif ret == -2:
        raise RuntimeError("Fused decode (zstd) failed (SHA256 mismatch)")
    elif ret != 0:
        raise RuntimeError(f"Fused decode (zstd) failed (ret={ret})")

    return output


def zstd_decompress(compressed: bytes, output_capacity: int) -> bytes:
    """
    Decompress zstd data using C++ kernel.
    WO-CDNA2-KERNEL-02: ~29% faster than brotli.

    Args:
        compressed: Zstd-compressed bytes
        output_capacity: Expected decompressed size

    Returns:
        Decompressed bytes

    Raises:
        RuntimeError: If decompression fails or kernel unavailable
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("C++ kernel not available")

    output = (c_uint8 * output_capacity)()
    output_len = c_size_t(0)

    compressed_arr = (c_uint8 * len(compressed)).from_buffer_copy(compressed)

    ret = lib.cdna_zstd_decompress(
        compressed_arr,
        len(compressed),
        output,
        output_capacity,
        byref(output_len),
    )

    if ret != 0:
        raise RuntimeError(f"Zstd decompression failed (ret={ret})")

    return bytes(output[:output_len.value])


# ============================================================
# WO-CDNA2-KERNEL-02 Phase 2: Fused Decode + Matmul
# ============================================================

def fused_decode_matmul(
    compressed_data: bytes,
    codebook: np.ndarray,
    X: np.ndarray,
    Y: np.ndarray,
    block_rows: int,
    cols: int,
    sidecar_positions: Optional[np.ndarray] = None,
    sidecar_values: Optional[np.ndarray] = None,
    accumulate: bool = True,
    codec: str = "brotli",
) -> None:
    """
    Fused decode + matmul: decompress → gather → sidecar → BLAS matmul.
    WO-CDNA2-KERNEL-02 Phase 2: Eliminates Python/C++ boundary crossing for matmul.

    Computes: Y += X @ W_block (if accumulate=True)
          or: Y = X @ W_block  (if accumulate=False)

    Where W_block [block_rows, cols] is decoded from compressed CDNA indices.

    Args:
        compressed_data: Compressed CDNA indices (brotli or zstd)
        codebook: float32 array of shape (256,)
        X: Input activations [batch_seq, block_rows], must be contiguous float32
        Y: Output accumulator [batch_seq, cols], modified in-place, must be contiguous float32
        block_rows: Number of rows in W_block (K dimension)
        cols: Number of columns in W_block (output dimension)
        sidecar_positions: Linear positions for corrections (optional)
        sidecar_values: Correction values (optional)
        accumulate: If True, Y += X@W; if False, Y = X@W
        codec: Compression codec - "brotli" or "zstd"

    Raises:
        RuntimeError: On decompression error or internal error
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("C++ kernel not available")

    codebook = np.ascontiguousarray(codebook, dtype=np.float32)
    if codebook.shape != (256,):
        raise ValueError(f"Codebook must have shape (256,), got {codebook.shape}")

    # Ensure X and Y are contiguous float32
    X = np.ascontiguousarray(X, dtype=np.float32)
    if not Y.flags['C_CONTIGUOUS']:
        raise ValueError("Y must be C-contiguous")
    if Y.dtype != np.float32:
        raise ValueError("Y must be float32")

    batch_seq = X.shape[0]

    # Prepare compressed data
    compressed_arr = (c_uint8 * len(compressed_data)).from_buffer_copy(compressed_data)

    # Prepare sidecar (nullable)
    if sidecar_positions is not None and sidecar_values is not None:
        sidecar_positions = np.ascontiguousarray(sidecar_positions, dtype=np.uint64)
        sidecar_values = np.ascontiguousarray(sidecar_values, dtype=np.float32)
        positions_ptr = sidecar_positions.ctypes.data_as(POINTER(c_uint64))
        values_ptr = sidecar_values.ctypes.data_as(POINTER(c_float))
        n_sidecar = len(sidecar_positions)
    else:
        positions_ptr = None
        values_ptr = None
        n_sidecar = 0

    # Get pointers
    codebook_ptr = codebook.ctypes.data_as(POINTER(c_float))
    X_ptr = X.ctypes.data_as(POINTER(c_float))
    Y_ptr = Y.ctypes.data_as(POINTER(c_float))

    # Choose codec
    if codec == "zstd":
        ret = lib.cdna_fused_decode_matmul_zstd(
            compressed_arr,
            len(compressed_data),
            codebook_ptr,
            X_ptr,
            Y_ptr,
            batch_seq,
            block_rows,
            cols,
            positions_ptr,
            values_ptr,
            n_sidecar,
            1 if accumulate else 0,
        )
    else:
        ret = lib.cdna_fused_decode_matmul(
            compressed_arr,
            len(compressed_data),
            codebook_ptr,
            X_ptr,
            Y_ptr,
            batch_seq,
            block_rows,
            cols,
            positions_ptr,
            values_ptr,
            n_sidecar,
            1 if accumulate else 0,
        )

    if ret != 0:
        raise RuntimeError(f"Fused decode+matmul failed (ret={ret})")


def fused_decode_matmul_transpose(
    compressed_data: bytes,
    codebook: np.ndarray,
    X: np.ndarray,
    Y: np.ndarray,
    block_rows: int,
    cols: int,
    sidecar_positions: Optional[np.ndarray] = None,
    sidecar_values: Optional[np.ndarray] = None,
    accumulate: bool = True,
    codec: str = "brotli",
) -> None:
    """
    Fused decode + matmul with transpose: Y += X @ W^T.
    WO-CDNA2-KERNEL-02 Phase 2: For GQA where K/V weights are stored transposed.

    Computes: Y += X @ W_block^T (if accumulate=True)
          or: Y = X @ W_block^T  (if accumulate=False)

    Where W_block is [block_rows, cols] and W_block^T is [cols, block_rows].
    X is [batch_seq, cols], Y is [batch_seq, block_rows].

    Args:
        compressed_data: Compressed CDNA indices
        codebook: float32 array of shape (256,)
        X: Input activations [batch_seq, cols]
        Y: Output accumulator [batch_seq, block_rows], modified in-place
        block_rows: W rows (becomes output cols after transpose)
        cols: W cols (input dimension after transpose)
        sidecar_positions: Linear positions for corrections (optional)
        sidecar_values: Correction values (optional)
        accumulate: If True, Y += X@W^T; if False, Y = X@W^T
        codec: Compression codec - "brotli" or "zstd"

    Raises:
        RuntimeError: On decompression error or internal error
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("C++ kernel not available")

    codebook = np.ascontiguousarray(codebook, dtype=np.float32)
    if codebook.shape != (256,):
        raise ValueError(f"Codebook must have shape (256,), got {codebook.shape}")

    # Ensure X and Y are contiguous float32
    X = np.ascontiguousarray(X, dtype=np.float32)
    if not Y.flags['C_CONTIGUOUS']:
        raise ValueError("Y must be C-contiguous")
    if Y.dtype != np.float32:
        raise ValueError("Y must be float32")

    batch_seq = X.shape[0]

    # Prepare compressed data
    compressed_arr = (c_uint8 * len(compressed_data)).from_buffer_copy(compressed_data)

    # Prepare sidecar (nullable)
    if sidecar_positions is not None and sidecar_values is not None:
        sidecar_positions = np.ascontiguousarray(sidecar_positions, dtype=np.uint64)
        sidecar_values = np.ascontiguousarray(sidecar_values, dtype=np.float32)
        positions_ptr = sidecar_positions.ctypes.data_as(POINTER(c_uint64))
        values_ptr = sidecar_values.ctypes.data_as(POINTER(c_float))
        n_sidecar = len(sidecar_positions)
    else:
        positions_ptr = None
        values_ptr = None
        n_sidecar = 0

    # Get pointers
    codebook_ptr = codebook.ctypes.data_as(POINTER(c_float))
    X_ptr = X.ctypes.data_as(POINTER(c_float))
    Y_ptr = Y.ctypes.data_as(POINTER(c_float))

    use_zstd = 1 if codec == "zstd" else 0

    ret = lib.cdna_fused_decode_matmul_transpose(
        compressed_arr,
        len(compressed_data),
        codebook_ptr,
        X_ptr,
        Y_ptr,
        batch_seq,
        block_rows,
        cols,
        positions_ptr,
        values_ptr,
        n_sidecar,
        1 if accumulate else 0,
        use_zstd,
    )

    if ret != 0:
        raise RuntimeError(f"Fused decode+matmul (transpose) failed (ret={ret})")


def compute_sha256(data: bytes) -> bytes:
    """Compute SHA256 using the C++ kernel."""
    lib = _load_library()
    if lib is None:
        raise RuntimeError("C++ kernel not available")

    data_arr = (c_uint8 * len(data)).from_buffer_copy(data)
    hash_out = (c_uint8 * 32)()

    ret = lib.cdna_sha256(data_arr, len(data), hash_out)
    if ret != 0:
        raise RuntimeError("SHA256 computation failed")

    return bytes(hash_out)


# ============================================================
# WO-CDNA2-ENCODER-KERNEL-01: Encoder Functions
# ============================================================

def get_encoder_version() -> Optional[str]:
    """Get the encoder kernel version string."""
    lib = _load_library()
    if lib is None:
        return None
    version_bytes = lib.cdna_encoder_version()
    return version_bytes.decode('utf-8') if version_bytes else None


def cpp_encoder_available() -> bool:
    """
    Check if the C++ encoder kernel is available AND enabled.

    The encoder is opt-in via HELIX_USE_CPP_ENCODER=1 environment variable.
    This allows gradual rollout and easy fallback to Python path.
    """
    if not _USE_CPP_ENCODER:
        return False
    return _load_library() is not None


def assign_nearest(
    values: np.ndarray,
    codebook: np.ndarray,
) -> np.ndarray:
    """
    Assign nearest centroid index for each value.
    Uses AVX2 SIMD acceleration when available.

    Deterministic tie-break: lowest index wins when equidistant.

    Args:
        values: float32 array of input values (any shape, will be flattened)
        codebook: float32 array of shape (256,) or smaller

    Returns:
        uint8 array of indices (same shape as values)

    Raises:
        RuntimeError: If kernel unavailable or call fails
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("C++ kernel not available")

    original_shape = values.shape

    # Ensure correct dtypes and contiguity
    values_flat = np.ascontiguousarray(values.flatten(), dtype=np.float32)
    codebook = np.ascontiguousarray(codebook, dtype=np.float32)

    # Pad codebook to 256 if smaller (k < 256 case)
    if len(codebook) < 256:
        padded_codebook = np.zeros(256, dtype=np.float32)
        padded_codebook[:len(codebook)] = codebook
        # Fill unused entries with +inf so they never win
        padded_codebook[len(codebook):] = np.finfo(np.float32).max
        codebook = padded_codebook
    elif len(codebook) > 256:
        raise ValueError(f"Codebook must have <= 256 entries, got {len(codebook)}")

    # Allocate output
    indices = np.empty(values_flat.shape, dtype=np.uint8)

    # Get ctypes pointers
    values_ptr = values_flat.ctypes.data_as(POINTER(c_float))
    codebook_ptr = codebook.ctypes.data_as(POINTER(c_float))
    indices_ptr = indices.ctypes.data_as(POINTER(c_uint8))

    ret = lib.cdna_assign_nearest(
        values_ptr,
        codebook_ptr,
        indices_ptr,
        len(values_flat),
    )

    if ret != 0:
        raise RuntimeError(f"Centroid assignment failed (ret={ret})")

    return indices.reshape(original_shape)


def assign_nearest_chunked(
    values: np.ndarray,
    codebook: np.ndarray,
    chunk_size: int = 500_000,
) -> np.ndarray:
    """
    Assign nearest centroid index with chunking for large arrays.

    Same as assign_nearest but processes in chunks to avoid memory issues.
    This matches the Python implementation's chunked approach.

    Args:
        values: float32 array of input values
        codebook: float32 array of shape (256,) or smaller
        chunk_size: Number of elements per chunk (default 500k)

    Returns:
        uint8 array of indices (same shape as values)
    """
    lib = _load_library()
    if lib is None:
        raise RuntimeError("C++ kernel not available")

    original_shape = values.shape
    flat = np.ascontiguousarray(values.flatten(), dtype=np.float32)
    codebook = np.ascontiguousarray(codebook, dtype=np.float32)

    # Pad codebook if needed
    if len(codebook) < 256:
        padded_codebook = np.zeros(256, dtype=np.float32)
        padded_codebook[:len(codebook)] = codebook
        padded_codebook[len(codebook):] = np.finfo(np.float32).max
        codebook = padded_codebook
    elif len(codebook) > 256:
        raise ValueError(f"Codebook must have <= 256 entries, got {len(codebook)}")

    indices = np.empty(len(flat), dtype=np.uint8)

    # Get codebook pointer (constant across chunks)
    codebook_ptr = codebook.ctypes.data_as(POINTER(c_float))

    for start in range(0, len(flat), chunk_size):
        end = min(start + chunk_size, len(flat))
        chunk = flat[start:end]
        chunk_indices = indices[start:end]

        chunk_ptr = chunk.ctypes.data_as(POINTER(c_float))
        indices_ptr = chunk_indices.ctypes.data_as(POINTER(c_uint8))

        ret = lib.cdna_assign_nearest(
            chunk_ptr,
            codebook_ptr,
            indices_ptr,
            len(chunk),
        )

        if ret != 0:
            raise RuntimeError(f"Centroid assignment failed (ret={ret}) at chunk {start}:{end}")

    return indices.reshape(original_shape)


def get_encoder_info() -> dict:
    """
    Get encoder kernel information for receipts.

    Returns:
        Dict with used, version, avx2 fields
    """
    return {
        "used": cpp_encoder_available(),
        "version": get_encoder_version(),
        "avx2": has_avx2(),
        "env_enabled": _USE_CPP_ENCODER,
    }


# Convenience: Check availability at import time
_available = cpp_kernel_available()
_encoder_available = cpp_encoder_available()
