"""
Cython extension for kHeavyHash using optimized C implementation.

This module provides a direct binding to the C kHeavyHash implementation,
offering significant performance improvements over the pure Python version.
"""

# cython: language_level=3
# cython: boundscheck=False
# cython: wraparound=False

from libc.stdint cimport uint8_t, uint16_t
from libc.stdlib cimport malloc, free
from libc.string cimport memcpy
from libc.stddef cimport size_t

# Declare the C function from kheavyhash_core.h
cdef extern from "kheavyhash_core.h":
    int kheavyhash_pow(const uint8_t *work_order80, uint8_t *output32)
    int kheavyhash_get_initial_hash(const uint8_t *work_order80, uint8_t *initial_hash32)
    int kheavyhash_get_xor_result(const uint8_t *work_order80, uint8_t *xor_result32)
    int kheavyhash_get_matrix(const uint8_t *work_order80, uint16_t *matrix_out)

# Declare cshake256 directly (it's in keccak_nano.c which is compiled with the extension)
cdef extern from "keccak_nano.h":
    int cshake256(uint8_t *out, const char *customization_str, const uint8_t *input_data, size_t input_len)


def kheavyhash(work_order: bytes) -> bytes:
    """
    kHeavyHash implementation using optimized C extension.

    This function directly calls the C implementation, providing
    significant performance improvements over the pure Python version.

    Args:
        work_order: Exactly 80 bytes of input data.
                   Format: [32 bytes PrePowHash][8 bytes Timestamp][32 bytes padding][8 bytes Nonce]

    Returns:
        32-byte hash digest.

    Raises:
        ValueError: If work_order is not exactly 80 bytes.
        RuntimeError: If the C implementation returns an error.
    """
    if len(work_order) != 80:
        raise ValueError(f"Work order must be exactly 80 bytes, got {len(work_order)}")

    # Prepare input buffer
    cdef const uint8_t[::1] work_order_view = work_order
    cdef uint8_t output[32]

    # Call the C implementation
    cdef int result = kheavyhash_pow(&work_order_view[0], output)

    if result != 0:
        raise RuntimeError("kHeavyHash C implementation returned error")

    # Convert output to Python bytes
    return bytes(output[:32])


def cshake256_hash(data: bytes, customization: str = "", output_len: int = 32) -> bytes:
    """
    cSHAKE256 hash using keccak_nano implementation (matching C code exactly).

    This function calls the C cshake256 function directly to ensure
    identical behavior to the C implementation.

    Args:
        data: Input data to hash
        customization: Customization string (default: empty string for SHAKE256)
        output_len: Desired output length in bytes (default: 32, max: 32)

    Returns:
        Hash digest of specified length (truncated from 32 bytes if output_len < 32)
    """
    if output_len > 32:
        raise ValueError("output_len cannot exceed 32 bytes")
    
    cdef const uint8_t[::1] data_view = data
    cdef uint8_t temp_output[32]
    cdef size_t data_len = len(data)
    cdef int result
    cdef bytes customization_bytes = customization.encode('utf-8')
    cdef const char *customization_cstr = customization_bytes
    
    # cshake256 always outputs 32 bytes
    result = cshake256(temp_output, customization_cstr, &data_view[0], data_len)
    if result != 0:
        raise RuntimeError("cshake256 returned error")
    
    # Return truncated output if requested
    return bytes(temp_output[:output_len])


def get_initial_hash(work_order: bytes) -> bytes:
    """
    Get the initial hash after ProofOfWorkHash step from C implementation.
    
    This exposes what the C implementation actually produces after the
    "ProofOfWorkHash" step, which uses keccak_nano cshake256.
    
    Args:
        work_order: Exactly 80 bytes of input data.
    
    Returns:
        32-byte initial hash from C implementation.
    
    Raises:
        ValueError: If work_order is not exactly 80 bytes.
        RuntimeError: If the C implementation returns an error.
    """
    if len(work_order) != 80:
        raise ValueError(f"Work order must be exactly 80 bytes, got {len(work_order)}")
    
    cdef const uint8_t[::1] work_order_view = work_order
    cdef uint8_t initial_hash[32]
    
    cdef int result = kheavyhash_get_initial_hash(&work_order_view[0], initial_hash)
    
    if result != 0:
        raise RuntimeError("kHeavyHash C implementation returned error")
    
    return bytes(initial_hash[:32])


def get_xor_result(work_order: bytes) -> bytes:
    """
    Get the XOR result after matrix operations (step 3 output) from C implementation.
    
    This exposes what the C implementation produces after:
    - Matrix-vector multiplication
    - XOR with original hash
    
    Args:
        work_order: Exactly 80 bytes of input data.
    
    Returns:
        32-byte XOR result from C implementation.
    
    Raises:
        ValueError: If work_order is not exactly 80 bytes.
        RuntimeError: If the C implementation returns an error.
    """
    if len(work_order) != 80:
        raise ValueError(f"Work order must be exactly 80 bytes, got {len(work_order)}")
    
    cdef const uint8_t[::1] work_order_view = work_order
    cdef uint8_t xor_result[32]
    
    cdef int result = kheavyhash_get_xor_result(&work_order_view[0], xor_result)
    
    if result != 0:
        raise RuntimeError("kHeavyHash C implementation returned error")
    
    return bytes(xor_result[:32])


def get_matrix(work_order: bytes) -> list:
    """
    Get the generated matrix (step 2 output) from C implementation.
    
    This exposes the 64x64 GF(16) matrix that was generated using xoshiro256++
    seeded from PrePowHash.
    
    Args:
        work_order: Exactly 80 bytes of input data.
    
    Returns:
        A 64x64 matrix as a list of lists (Python ints).
    
    Raises:
        ValueError: If work_order is not exactly 80 bytes.
        RuntimeError: If the C implementation returns an error.
    """
    if len(work_order) != 80:
        raise ValueError(f"Work order must be exactly 80 bytes, got {len(work_order)}")
    
    cdef const uint8_t[::1] work_order_view = work_order
    cdef uint16_t matrix_flat[64 * 64]
    
    cdef int result = kheavyhash_get_matrix(&work_order_view[0], matrix_flat)
    
    if result != 0:
        raise RuntimeError("kHeavyHash C implementation returned error")
    
    # Convert flat array to 64x64 matrix
    matrix = []
    for i in range(64):
        row = []
        for j in range(64):
            row.append(int(matrix_flat[i * 64 + j]))
        matrix.append(row)
    
    return matrix
