#include "kheavyhash_core.h"
#include "keccak_nano.h"

#include <math.h>
#include <string.h>

#define KHEAVYHASH_EPS 1e-9

static inline uint64_t kheavyhash_rotl64(const uint64_t x, int k) {
    return (x << k) | (x >> (64 - k));
}

uint64_t kheavyhash_xoshiro_next(struct kheavyhash_xoshiro_state *state) {
    const uint64_t result = kheavyhash_rotl64(state->s[0] + state->s[3], 23) + state->s[0];

    const uint64_t t = state->s[1] << 17;

    state->s[2] ^= state->s[0];
    state->s[3] ^= state->s[1];
    state->s[1] ^= state->s[2];
    state->s[0] ^= state->s[3];

    state->s[2] ^= t;

    state->s[3] = kheavyhash_rotl64(state->s[3], 45);

    return result;
}

int kheavyhash_compute_rank(const uint16_t matrix[KHEAVYHASH_MATRIX_SIZE][KHEAVYHASH_MATRIX_SIZE]) {
    /*
     * The working matrix must be double. Gaussian elimination normalizes the
     * pivot row by division and then subtracts multiples to eliminate entries.
     * With integer types, division truncates (e.g. 7/3 = 2); the subsequent
     * elimination then fails to cancel exactly, leaving spurious non-zeros and
     * wrong rank. Doubles avoid that. There are other ways to do rank over
     * this input (e.g. proper GF(16) elimination) that avoid floats entirely,
     * but this was simpler to implement, and this path is not performance-
     * critical, so it is fine as-is.
     */
    double B[KHEAVYHASH_MATRIX_SIZE][KHEAVYHASH_MATRIX_SIZE];
    bool row_selected[KHEAVYHASH_MATRIX_SIZE] = {0};
    int rank = 0;

    for (int i = 0; i < KHEAVYHASH_MATRIX_SIZE; ++i) {
        for (int j = 0; j < KHEAVYHASH_MATRIX_SIZE; ++j) {
            B[i][j] = (double)matrix[i][j];
        }
    }

    for (int col = 0; col < KHEAVYHASH_MATRIX_SIZE; ++col) {
        int pivot_row = -1;
        for (int row = 0; row < KHEAVYHASH_MATRIX_SIZE; ++row) {
            if (!row_selected[row] && fabs(B[row][col]) > KHEAVYHASH_EPS) {
                pivot_row = row;
                break;
            }
        }
        if (pivot_row == -1) {
            continue;
        }

        ++rank;
        row_selected[pivot_row] = true;

        /* Normalize pivot row */
        for (int p = col + 1; p < KHEAVYHASH_MATRIX_SIZE; ++p) {
            B[pivot_row][p] /= B[pivot_row][col];
        }

        /* Eliminate column in other rows */
        for (int k = 0; k < KHEAVYHASH_MATRIX_SIZE; ++k) {
            if (k != pivot_row && fabs(B[k][col]) > KHEAVYHASH_EPS) {
                for (int p = col + 1; p < KHEAVYHASH_MATRIX_SIZE; ++p) {
                    B[k][p] -= B[pivot_row][p] * B[k][col];
                }
            }
        }
    }

    return rank;
}

bool kheavyhash_is_full_rank(
    const uint16_t matrix[KHEAVYHASH_MATRIX_SIZE][KHEAVYHASH_MATRIX_SIZE]
) {
    return kheavyhash_compute_rank(matrix) == KHEAVYHASH_MATRIX_SIZE;
}

void kheavyhash_generate_matrix(
    uint16_t matrix[KHEAVYHASH_MATRIX_SIZE][KHEAVYHASH_MATRIX_SIZE],
    struct kheavyhash_xoshiro_state *state
) {
    do {
        for (int i = 0; i < KHEAVYHASH_MATRIX_SIZE; ++i) {
            for (int j = 0; j < KHEAVYHASH_MATRIX_SIZE; j += 16) {
                uint64_t value = kheavyhash_xoshiro_next(state);
                for (int shift = 0; shift < 16; ++shift) {
                    matrix[i][j + shift] = (uint16_t)((value >> (4 * shift)) & 0xF);
                }
            }
        }
    } while (!kheavyhash_is_full_rank(matrix));
}

void kheavyhash_heavyhash(
    const uint16_t matrix[KHEAVYHASH_MATRIX_SIZE][KHEAVYHASH_MATRIX_SIZE],
    const uint8_t *seed32,
    uint8_t *output32
) {
    uint8_t hash_first[32] __attribute__((aligned(64)));
    uint8_t hash_second[32] __attribute__((aligned(64)));

    uint16_t vector[64] __attribute__((aligned(64)));
    uint16_t product[64] __attribute__((aligned(64)));

    /* First hash input is provided as seed32 */
    memcpy(hash_first, seed32, 32);

    /* Expand bytes into 64 GF(16) elements (nibbles) */
    for (int i = 0; i < 32; ++i) {
        vector[2 * i] = (uint16_t)(hash_first[i] >> 4);
        vector[2 * i + 1] = (uint16_t)(hash_first[i] & 0xF);
    }

    /* Matrix-vector multiplication with scaling */
    for (int i = 0; i < 64; ++i) {
        uint16_t sum = 0;
        for (int j = 0; j < 64; ++j) {
            sum = (uint16_t)(sum + matrix[i][j] * vector[j]);
        }
        product[i] = (uint16_t)(sum >> KHEAVYHASH_MATRIX_VECTOR_SCALE_SHIFT);
    }

    /* Recombine nibbles into 32 bytes */
    for (int i = 0; i < 32; ++i) {
        hash_second[i] = (uint8_t)((product[2 * i] << 4) | (product[2 * i + 1]));
    }

    /* XOR with original - output is the XOR result */
    for (int i = 0; i < 32; ++i) {
        output32[i] = (uint8_t)(hash_first[i] ^ hash_second[i]);
    }
}


int kheavyhash_pow(
    const uint8_t *work_order80,
    uint8_t *output32
) {
    if (work_order80 == NULL || output32 == NULL) {
        return -1;
    }

    uint8_t initial_hash[32] = {0};
    uint16_t matrix[KHEAVYHASH_MATRIX_SIZE][KHEAVYHASH_MATRIX_SIZE];
    struct kheavyhash_xoshiro_state state;

    /*
     * Step 1: Hash the 80-byte work order with cSHAKE256("ProofOfWorkHash")
     *
     * The work order format:
     *   - PrePowHash (32 bytes, bytes 0-31)
     *   - Timestamp (8 bytes, bytes 32-39)
     *   - Padding (32 bytes of zeros, bytes 40-71) - included in hash
     *   - Nonce (8 bytes, bytes 72-79)
     */
    if (cshake256(initial_hash, "ProofOfWorkHash", work_order80, 80) != 0) {
        return -1;
    }

    /*
     * Step 2: Seed xoshiro256++ PRNG from PrePowHash (first 32 bytes of work_order)
     *
     * CRITICAL: The PRNG seed comes from PrePowHash, NOT from the initial hash!
     * This matches the original miner code: state.s[i] = ((uint64_t*)edata)[i]
     * where edata contains the PrePowHash bytes.
     */
    const uint8_t *pre_pow_hash = work_order80;
    for (int i = 0; i < 4; ++i) {
        /* Extract 4 uint64 values from PrePowHash (little-endian) */
        state.s[i] = ((uint64_t)pre_pow_hash[i * 8 + 0]) |
                     ((uint64_t)pre_pow_hash[i * 8 + 1] << 8) |
                     ((uint64_t)pre_pow_hash[i * 8 + 2] << 16) |
                     ((uint64_t)pre_pow_hash[i * 8 + 3] << 24) |
                     ((uint64_t)pre_pow_hash[i * 8 + 4] << 32) |
                     ((uint64_t)pre_pow_hash[i * 8 + 5] << 40) |
                     ((uint64_t)pre_pow_hash[i * 8 + 6] << 48) |
                     ((uint64_t)pre_pow_hash[i * 8 + 7] << 56);
    }

    /*
     * Step 3: Generate a full-rank 64x64 GF(16) matrix
     */
    kheavyhash_generate_matrix(matrix, &state);

    /*
     * Step 4: Apply the heavyhash transformation to the initial hash
     *
     * The heavyhash function takes the 32-byte initial hash, performs
     * matrix-vector multiplication, and XORs with the original.
     */
    uint8_t hash_xored[32] __attribute__((aligned(64)));
    kheavyhash_heavyhash(matrix, initial_hash, hash_xored);

    /*
     * Step 5: Final hash with cSHAKE256("HeavyHash")
     */
    if (cshake256(output32, "HeavyHash", hash_xored, 32) != 0) {
        return -1;
    }

    return 0;
}

int kheavyhash_get_initial_hash(
    const uint8_t *work_order80,
    uint8_t *initial_hash32
) {
    if (work_order80 == NULL || initial_hash32 == NULL) {
        return -1;
    }

    /*
     * Step 1: Hash the 80-byte work order with cSHAKE256("ProofOfWorkHash")
     */
    if (cshake256(initial_hash32, "ProofOfWorkHash", work_order80, 80) != 0) {
        return -1;
    }

    return 0;
}

int kheavyhash_get_xor_result(
    const uint8_t *work_order80,
    uint8_t *xor_result32
) {
    if (work_order80 == NULL || xor_result32 == NULL) {
        return -1;
    }

    uint8_t initial_hash[32] = {0};
    uint16_t matrix[KHEAVYHASH_MATRIX_SIZE][KHEAVYHASH_MATRIX_SIZE];
    struct kheavyhash_xoshiro_state state;

    /* Step 1: Initial hash */
    if (cshake256(initial_hash, "ProofOfWorkHash", work_order80, 80) != 0) {
        return -1;
    }

    /* Step 2: Seed PRNG from PrePowHash */
    const uint8_t *pre_pow_hash = work_order80;
    for (int i = 0; i < 4; ++i) {
        state.s[i] = ((uint64_t)pre_pow_hash[i * 8 + 0]) |
                     ((uint64_t)pre_pow_hash[i * 8 + 1] << 8) |
                     ((uint64_t)pre_pow_hash[i * 8 + 2] << 16) |
                     ((uint64_t)pre_pow_hash[i * 8 + 3] << 24) |
                     ((uint64_t)pre_pow_hash[i * 8 + 4] << 32) |
                     ((uint64_t)pre_pow_hash[i * 8 + 5] << 40) |
                     ((uint64_t)pre_pow_hash[i * 8 + 6] << 48) |
                     ((uint64_t)pre_pow_hash[i * 8 + 7] << 56);
    }

    /* Step 3: Generate matrix */
    kheavyhash_generate_matrix(matrix, &state);

    /* Step 4: Matrix operations and XOR */
    kheavyhash_heavyhash(matrix, initial_hash, xor_result32);

    return 0;
}

int kheavyhash_get_matrix(
    const uint8_t *work_order80,
    uint16_t *matrix_out
) {
    if (work_order80 == NULL || matrix_out == NULL) {
        return -1;
    }

    uint16_t matrix[KHEAVYHASH_MATRIX_SIZE][KHEAVYHASH_MATRIX_SIZE];
    struct kheavyhash_xoshiro_state state;

    /* Seed PRNG from PrePowHash */
    const uint8_t *pre_pow_hash = work_order80;
    for (int i = 0; i < 4; ++i) {
        state.s[i] = ((uint64_t)pre_pow_hash[i * 8 + 0]) |
                     ((uint64_t)pre_pow_hash[i * 8 + 1] << 8) |
                     ((uint64_t)pre_pow_hash[i * 8 + 2] << 16) |
                     ((uint64_t)pre_pow_hash[i * 8 + 3] << 24) |
                     ((uint64_t)pre_pow_hash[i * 8 + 4] << 32) |
                     ((uint64_t)pre_pow_hash[i * 8 + 5] << 40) |
                     ((uint64_t)pre_pow_hash[i * 8 + 6] << 48) |
                     ((uint64_t)pre_pow_hash[i * 8 + 7] << 56);
    }

    /* Generate matrix */
    kheavyhash_generate_matrix(matrix, &state);

    /* Copy matrix to output (row-major order) */
    for (int i = 0; i < KHEAVYHASH_MATRIX_SIZE; ++i) {
        for (int j = 0; j < KHEAVYHASH_MATRIX_SIZE; ++j) {
            matrix_out[i * KHEAVYHASH_MATRIX_SIZE + j] = matrix[i][j];
        }
    }

    return 0;
}

