"""
Pure Python implementation of kHeavyHash.

This module provides a 100% pure Python implementation of the kHeavyHash
algorithm using standard libraries and pycryptodome for cSHAKE256 support.
"""

import struct
from typing import List, Tuple, cast

try:
    from Crypto.Hash import cSHAKE256
except ImportError:
    raise ImportError(
        "pycryptodome is required for cSHAKE256 support. "
        "Install it with: pip install pycryptodome"
    )

# Matrix dimensions
MATRIX_SIZE = 64
GF16_MODULUS = 16

# Maximum attempts to generate a full-rank matrix
# Most random matrices over GF(16) are full-rank, but some seeds may require
# thousands of attempts. This limit ensures we don't hang indefinitely.
MAX_MATRIX_GENERATION_ATTEMPTS = 10000

# Bit shift amount for scaling matrix-vector product results
# This reduces the magnitude of the dot product to fit in GF(16) range
MATRIX_VECTOR_SCALE_SHIFT = 10

# Epsilon for floating-point rank computation (matching C implementation)
EPS = 1e-9


class Xoshiro256PlusPlus:
    """
    xoshiro256++ pseudorandom number generator.

    This is a fast, high-quality PRNG used for matrix generation in kHeavyHash.
    The implementation matches the C and Go reference implementations exactly.

    CRITICAL: Python integers are arbitrary precision and don't wrap like C uint64_t.
    All state operations must be explicitly masked to 64 bits to match C behavior.

    Per Go reference: the 4 uint64 seed values are used directly as the PRNG state.

    Args:
        seed: Four 64-bit unsigned integers (little-endian) as the initial state.
    """

    def __init__(self, seed: Tuple[int, int, int, int]):
        """
        Initialize xoshiro256++ with a 256-bit seed.

        Per Go reference: the seed values are used directly as the PRNG state.

        Args:
            seed: Tuple of four 64-bit unsigned integers representing the state.
        """
        if len(seed) != 4:
            raise ValueError("Seed must contain exactly 4 uint64 values")
        # Use seed values directly as state (per Go reference implementation)
        self.state = list(seed)

    def _rotl(self, x: int, k: int) -> int:
        """Rotate left by k bits (64-bit)."""
        x = x & 0xFFFFFFFFFFFFFFFF  # Ensure 64-bit
        return ((x << k) | (x >> (64 - k))) & 0xFFFFFFFFFFFFFFFF

    def next(self) -> int:
        """
        Generate the next 64-bit random number.

        This matches the C implementation exactly:
        - result = rotl(s[0] + s[3], 23) + s[0]
        - t = s[1] << 17
        - s[2] ^= s[0]
        - s[3] ^= s[1]
        - s[1] ^= s[2]
        - s[0] ^= s[3]
        - s[2] ^= t
        - s[3] = rotl(s[3], 45)

        Returns:
            A 64-bit unsigned integer.
        """
        # Ensure state values are 64-bit
        s0 = self.state[0] & 0xFFFFFFFFFFFFFFFF
        s1 = self.state[1] & 0xFFFFFFFFFFFFFFFF
        s2 = self.state[2] & 0xFFFFFFFFFFFFFFFF
        s3 = self.state[3] & 0xFFFFFFFFFFFFFFFF

        # result = rotl(s[0] + s[3], 23) + s[0]
        result = (self._rotl((s0 + s3) & 0xFFFFFFFFFFFFFFFF, 23) + s0) & 0xFFFFFFFFFFFFFFFF

        # t = s[1] << 17
        t = (s1 << 17) & 0xFFFFFFFFFFFFFFFF

        # Update state (must mask after each operation to match C uint64_t behavior)
        s2 = (s2 ^ s0) & 0xFFFFFFFFFFFFFFFF
        s3 = (s3 ^ s1) & 0xFFFFFFFFFFFFFFFF
        s1 = (s1 ^ s2) & 0xFFFFFFFFFFFFFFFF
        s0 = (s0 ^ s3) & 0xFFFFFFFFFFFFFFFF

        s2 = (s2 ^ t) & 0xFFFFFFFFFFFFFFFF

        s3 = self._rotl(s3, 45)

        # Update state
        self.state[0] = s0
        self.state[1] = s1
        self.state[2] = s2
        self.state[3] = s3

        return result


def cshake256(data: bytes, domain: str, output_len: int = 32) -> bytes:
    """
    Compute cSHAKE256 hash with domain separation using pycryptodome.

    Args:
        data: Input data to hash.
        domain: Customization string for domain separation.
        output_len: Desired output length in bytes (default: 32).

    Returns:
        Hash digest of specified length.

    Example:
        >>> result = cshake256(b"hello", "test", output_len=32)
        >>> len(result)
        32
    """
    shake = cSHAKE256.new(data=data, custom=domain.encode("utf-8"))
    return cast(bytes, shake.read(output_len))


def matrix_rank_float(matrix: List[List[int]]) -> int:
    """
    Compute the rank of a matrix using floating-point arithmetic.

    This exactly matches the C implementation's kheavyhash_compute_rank() which uses
    floating-point (double) arithmetic.

    Args:
        matrix: Square matrix of integers (treated as floats for rank computation)

    Returns:
        The rank of the matrix (0 to MATRIX_SIZE)
    """
    n = len(matrix)
    # Convert to float matrix (matching C: double b[64][64])
    b = [[float(matrix[i][j]) for j in range(n)] for i in range(n)]  # noqa: N806
    row_selected = [False] * n
    rank = 0

    for col in range(n):
        # Find pivot row (matching C: first row with non-zero in this column)
        pivot_row = None
        for row in range(n):
            if not row_selected[row] and abs(b[row][col]) > EPS:
                pivot_row = row
                break

        if pivot_row is None:
            continue

        rank += 1
        row_selected[pivot_row] = True

        # Normalize pivot row (matching C: b[pivot_row][p] /= b[pivot_row][col])
        # Only normalize columns after the pivot column
        for p in range(col + 1, n):
            b[pivot_row][p] /= b[pivot_row][col]

        # Eliminate column in other rows (matching C implementation exactly)
        for k in range(n):
            if k != pivot_row and abs(b[k][col]) > EPS:
                for p in range(col + 1, n):
                    b[k][p] -= b[pivot_row][p] * b[k][col]

    return rank


def matrix_rank_gf16(matrix: List[List[int]]) -> int:
    """
    Compute the rank of a matrix over GF(16).

    Uses Gaussian elimination to determine the rank. The matrix must be square
    and contain values in the range [0, 15] (GF(16) elements).

    Args:
        matrix: Square matrix of GF(16) elements (nibbles 0-15).

    Returns:
        The rank of the matrix (0 to MATRIX_SIZE).

    Example:
        >>> identity = [[1 if i == j else 0 for j in range(64)] for i in range(64)]
        >>> matrix_rank_gf16(identity)
        64
    """
    n = len(matrix)
    # Create a copy to avoid modifying the original
    m = [row[:] for row in matrix]
    rank = 0

    for col in range(n):
        # Find pivot row
        pivot_row = None
        for row in range(rank, n):
            if m[row][col] != 0:
                pivot_row = row
                break

        if pivot_row is None:
            continue

        # Swap rows
        if pivot_row != rank:
            m[rank], m[pivot_row] = m[pivot_row], m[rank]

        # Normalize pivot row (make pivot element 1 in GF(16))
        pivot_val = m[rank][col]
        if pivot_val != 1:
            # Find multiplicative inverse in GF(16)
            inv = _gf16_inverse(pivot_val)
            for c in range(n):
                m[rank][c] = (m[rank][c] * inv) % GF16_MODULUS

        # Eliminate column
        for row in range(n):
            if row != rank and m[row][col] != 0:
                factor = m[row][col]
                for c in range(n):
                    m[row][c] = (m[row][c] - factor * m[rank][c]) % GF16_MODULUS

        rank += 1

    return rank


def _gf16_inverse(a: int) -> int:
    """
    Compute the multiplicative inverse in GF(16).

    Uses extended Euclidean algorithm. GF(16) is isomorphic to GF(2)[x]/(x^4 + x + 1).

    Args:
        a: Element in GF(16) (0-15). Must not be 0.

    Returns:
        Multiplicative inverse of a in GF(16).

    Raises:
        ValueError: If a is 0 (no inverse exists).
    """
    if a == 0:
        raise ValueError("Cannot compute inverse of 0 in GF(16)")

    # Precomputed inverses for GF(16) using irreducible polynomial x^4 + x + 1
    # This is faster than computing on the fly
    inverses = [0, 1, 9, 14, 13, 11, 7, 6, 15, 2, 5, 10, 12, 3, 4, 8]
    return inverses[a]


def generate_matrix(prng: Xoshiro256PlusPlus) -> List[List[int]]:
    """
    Generate a 64x64 matrix over GF(16) using the PRNG.

    This matches kheavyhash_generate_matrix() in the C implementation.
    The C code has a do-while loop INSIDE the function that regenerates the matrix until full rank.
    This function must also have the do-while loop inside to match PRNG state consumption.

    Each matrix element is a 4-bit nibble (0-15). The matrix is filled row by row,
    with each uint64 from the PRNG providing 16 matrix elements (4 bits each).

    Args:
        prng: Initialized xoshiro256++ PRNG instance.

    Returns:
        A 64x64 matrix of GF(16) elements (nibbles 0-15) that is guaranteed to be full rank.

    Example:
        >>> seed = (0x1234567890ABCDEF, 0xFEDCBA0987654321, 0x1111222233334444, 0x5555666677778888)
        >>> prng = Xoshiro256PlusPlus(seed)
        >>> matrix = generate_matrix(prng)
        >>> len(matrix)
        64
        >>> len(matrix[0])
        64
    """
    # CRITICAL: The do-while loop must be INSIDE this function to match C implementation.
    # The C code has: do { ... generate matrix ... } while (!is_full_rank(matrix));
    # This ensures PRNG state consumption matches exactly.
    matrix = None
    for attempt in range(MAX_MATRIX_GENERATION_ATTEMPTS):
        matrix = []
        for i in range(MATRIX_SIZE):
            row = []
            for j in range(0, MATRIX_SIZE, 16):
                # Get next uint64 from PRNG
                val = prng.next()
                # Extract 16 nibbles (4 bits each) from the uint64
                # Matching C: matrix[i][j + shift] = (uint16_t)((value >> (4 * shift)) & 0xF);
                for shift in range(16):
                    nibble = (val >> (4 * shift)) & 0xF
                    row.append(nibble)
            matrix.append(row)

        # Check if matrix is full rank (matching C: while (!is_full_rank(matrix))
        rank = matrix_rank_float(matrix)
        if rank == MATRIX_SIZE:
            break
    else:
        raise RuntimeError(
            f"Failed to generate full-rank matrix after {MAX_MATRIX_GENERATION_ATTEMPTS} attempts"
        )

    return matrix


def bytes_to_nibbles(data: bytes) -> List[int]:
    """
    Convert bytes to a list of nibbles (4-bit values).

    Each byte is split into two nibbles: upper 4 bits (high nibble) first,
    then lower 4 bits (low nibble).

    Args:
        data: Input bytes.

    Returns:
        List of nibbles (4-bit values, 0-15).

    Raises:
        TypeError: If data is not bytes.

    Example:
        >>> bytes_to_nibbles(b'\\x12\\x34')
        [1, 2, 3, 4]
    """
    if not isinstance(data, bytes):
        raise TypeError(f"data must be bytes, got {type(data).__name__}")
    nibbles = []
    for byte_val in data:
        # High nibble (upper 4 bits) first, then low nibble (lower 4 bits)
        nibbles.append((byte_val >> 4) & 0xF)
        nibbles.append(byte_val & 0xF)
    return nibbles


def nibbles_to_bytes(nibbles: List[int]) -> bytes:
    """
    Convert a list of nibbles to bytes.

    Pairs of nibbles are combined into bytes: (nibble[2*i] << 4) | nibble[2*i+1].

    Args:
        nibbles: List of nibbles (4-bit values, 0-15). Length must be even.

    Returns:
        Bytes representation of the nibbles.

    Raises:
        ValueError: If length is odd or any nibble is outside range [0, 15].

    Example:
        >>> nibbles_to_bytes([1, 2, 3, 4])
        b'\\x12\\x34'
    """
    if len(nibbles) % 2 != 0:
        raise ValueError("Nibbles list must have even length")
    result = []
    for i in range(0, len(nibbles), 2):
        # Matching C exactly:
        # hash_second[i] = (uint8_t)((product[2 * i] << 4) | (product[2 * i + 1]));
        # C code doesn't mask - it relies on product values being <= 15 after right-shift
        # We must match C exactly, so no masking
        byte_val = (nibbles[i] << 4) | nibbles[i + 1]
        result.append(byte_val)
    return bytes(result)


def matrix_vector_multiply_gf16(matrix: List[List[int]], vector: List[int]) -> List[int]:
    """
    Multiply a matrix by a vector over GF(16).

    Computes: result[i] = sum(matrix[i][j] * vector[j]) for each row i.
    After computing the sum, the result is right-shifted by 10 bits to scale it.

    CRITICAL: The C implementation uses uint16_t for the sum, which wraps around
    on overflow. Python ints are arbitrary precision, so we must explicitly
    mask to uint16_t range (0-65535) to match C behavior.

    Args:
        matrix: 64x64 matrix of GF(16) elements.
        vector: Vector of 64 GF(16) elements.

    Returns:
        Result vector of 64 elements (after scaling).
    """
    if len(matrix) != MATRIX_SIZE or len(vector) != MATRIX_SIZE:
        raise ValueError(f"Matrix and vector must have size {MATRIX_SIZE}")

    result = []
    for i in range(MATRIX_SIZE):
        # Compute dot product - CRITICAL: match C's uint16_t wrapping behavior
        # C code: uint16_t sum = 0; sum = (uint16_t)(sum + matrix[i][j] * vector[j]);
        # This means sum wraps around at 65536 (0x10000)
        dot_product = 0
        for j in range(MATRIX_SIZE):
            # Explicitly wrap to uint16_t range to match C behavior
            dot_product = (dot_product + matrix[i][j] * vector[j]) & 0xFFFF
        # Right-shift to scale (matching C: product[i] = (uint16_t)(sum >> 10))
        scaled = dot_product >> MATRIX_VECTOR_SCALE_SHIFT
        result.append(scaled)

    return result


def kheavyhash(work_order: bytes) -> bytes:
    """
    Compute kHeavyHash Proof-of-Work hash.

    This is the main entry point for the kHeavyHash algorithm. It implements
    the "Keccak-Matrix-Keccak" sandwich:

    1. Hash the 80-byte work order with cSHAKE256("ProofOfWorkHash")
    2. Seed xoshiro256++ PRNG from PrePowHash (first 32 bytes of work_order)
    3. Generate a 64x64 matrix over GF(16) using the PRNG, ensuring full-rank
    4. Extract 64 nibbles from the hash to form a vector
    5. Multiply matrix by vector over GF(16), scale by right-shifting 10 bits
    6. Reconstruct 32 bytes from the result vector
    7. XOR with original hash
    8. Hash with cSHAKE256("HeavyHash")
    9. Return the final hash (no byte reversal)

    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.
    """

    if len(work_order) != 80:
        raise ValueError(f"Work order must be exactly 80 bytes, got {len(work_order)}")

    # Step 1: First hash with cSHAKE256("ProofOfWorkHash")
    hashed_header = cshake256(work_order, "ProofOfWorkHash", output_len=32)

    # Step 2: Seed xoshiro256++ PRNG from PrePowHash
    # (first 32 bytes of input), NOT from hashed_header!
    # According to Go reference: prePowHash := newDomainHashFromByteArray((*[32]byte)(input[:32]))
    pre_pow_hash = work_order[:32]
    if len(pre_pow_hash) != 32:
        raise ValueError(f"pre_pow_hash must be exactly 32 bytes, got {len(pre_pow_hash)}")
    seed = struct.unpack("<4Q", pre_pow_hash)  # 4 unsigned long long (uint64) from PrePowHash

    # Step 3: Generate full-rank matrix
    # Matching C: kheavyhash_generate_matrix() has do-while loop inside
    # The C code uses floating-point rank computation, so we use that too
    prng = Xoshiro256PlusPlus(seed)
    # The do-while loop is now INSIDE generate_matrix() to match C implementation
    matrix = generate_matrix(prng)

    # Step 4: Extract vector from hash (64 nibbles from 32 bytes)
    vector = bytes_to_nibbles(hashed_header)

    # Step 5: Matrix-vector multiplication over GF(16)
    product_vector = matrix_vector_multiply_gf16(matrix, vector)

    # Step 6: Reconstruct 32 bytes from product vector
    product_bytes = nibbles_to_bytes(product_vector)

    # Step 7: XOR with original hash
    xor_result = bytes(a ^ b for a, b in zip(hashed_header, product_bytes))

    # Step 8: Final hash with cSHAKE256("HeavyHash")
    final_hash = cshake256(xor_result, "HeavyHash", output_len=32)

    # Step 9: Return final hash (per Go reference: return *heavyHash.byteArray() - NO reversal)
    # Go code shows: return *heavyHash.byteArray() - no byte reversal
    return final_hash
