/*
 * Copyright (c) 2025, magnum
 * This software is hereby released to the general public under
 * the following terms: Redistribution and use in source and binary
 * forms, with or without modification, are permitted.
 */

#include "opencl_device_info.h"
#define AMD_PUTCHAR_NOCAST
#include "opencl_misc.h"
#include "opencl_mask.h"
#include "opencl_sha2_ctx.h"

#define SWAP32(n) (((n) << 24) | (((n) & 0xff00U) << 8) | (((n) >> 8) & 0xff00U) | ((n) >> 24))

/* This handles an input of 0xffffffffU correctly */
#define BITMAP_SHIFT ((BITMAP_MASK >> 5) + 1)

#define sha256_single(W, out) { \
        uint ctx[8]; \
        sha256_init(ctx); \
        sha256_block(W, ctx); \
        for (int _i = 0; _i < 8; _i++) \
            (out)[_i] = ctx[_i]; \
    }

INLINE void cmp_final(uint gid,
                      uint iter,
                      uint *hash,
                      __global uint *offset_table,
                      __global uint *hash_table,
                      __global uint *return_hashes,
                      volatile __global uint *output,
                      volatile __global uint *bitmap_dupe)
{

	uint t, offset_table_index, hash_table_index;
	ulong LO, MI, HI;
	ulong p;

	HI = ((ulong)hash[5] << 32) | (ulong)hash[4];
	MI = ((ulong)hash[3] << 32) | (ulong)hash[2];
	LO = ((ulong)hash[1] << 32) | (ulong)hash[0];

	p = (HI % OFFSET_TABLE_SIZE) * SHIFT128_OT_SZ;
	p += (MI % OFFSET_TABLE_SIZE) * SHIFT64_OT_SZ;
	p += LO % OFFSET_TABLE_SIZE;
	p %= OFFSET_TABLE_SIZE;
	offset_table_index = (uint)p;

	//error: chances of overflow is extremely low.
	LO += (ulong)offset_table[offset_table_index];

	p = (HI % HASH_TABLE_SIZE) * SHIFT128_HT_SZ;
	p += (MI % HASH_TABLE_SIZE) * SHIFT64_HT_SZ;
	p += LO % HASH_TABLE_SIZE;
	p %= HASH_TABLE_SIZE;
	hash_table_index = (uint)p;

	if (hash_table[hash_table_index] == hash[0])
		if (hash_table[HASH_TABLE_SIZE + hash_table_index] == hash[1]) {
			/*
			 * Prevent duplicate keys from cracking same hash
			 */
			if (!(atomic_or(&bitmap_dupe[hash_table_index / 32],
			                (1U << (hash_table_index % 32))) & (1U << (hash_table_index % 32)))) {
				t = atomic_inc(&output[0]);
				output[1 + 3 * t] = gid;
				output[2 + 3 * t] = iter;
				output[3 + 3 * t] = hash_table_index;
				return_hashes[2 * t] = hash[6];
				return_hashes[2 * t + 1] = hash[7];
			}
		}
}

INLINE void cmp(uint gid,
                uint iter,
                uint *hash,
#if USE_LOCAL_BITMAPS
                __local
#else
                __global
#endif
                uint *bitmaps,
                __global uint *offset_table,
                __global uint *hash_table,
                __global uint *return_hashes,
                volatile __global uint *output,
                volatile __global uint *bitmap_dupe)
{
	uint bitmap_index, tmp = 1;

#if SELECT_CMP_STEPS > 4
	bitmap_index = hash[0] & BITMAP_MASK;
	tmp &= (bitmaps[bitmap_index >> 5] >> (bitmap_index & 31)) & 1U;
	bitmap_index = (hash[0] >> 16) & BITMAP_MASK;
	tmp &= (bitmaps[BITMAP_SHIFT + (bitmap_index >> 5)] >> (bitmap_index & 31)) &
	       1U;
	bitmap_index = hash[1] & BITMAP_MASK;
	tmp &= (bitmaps[BITMAP_SHIFT * 2 + (bitmap_index >> 5)] >>
	        (bitmap_index & 31)) & 1U;
	bitmap_index = (hash[1] >> 16) & BITMAP_MASK;
	tmp &= (bitmaps[BITMAP_SHIFT * 3 + (bitmap_index >> 5)] >>
	        (bitmap_index & 31)) & 1U;
	bitmap_index = hash[2] & BITMAP_MASK;
	tmp &= (bitmaps[BITMAP_SHIFT * 4 + (bitmap_index >> 5)] >>
	        (bitmap_index & 31)) & 1U;
	bitmap_index = (hash[2] >> 16) & BITMAP_MASK;
	tmp &= (bitmaps[BITMAP_SHIFT * 5 + (bitmap_index >> 5)] >>
	        (bitmap_index & 31)) & 1U;
	bitmap_index = hash[3] & BITMAP_MASK;
	tmp &= (bitmaps[BITMAP_SHIFT * 6 + (bitmap_index >> 5)] >>
	        (bitmap_index & 31)) & 1U;
	bitmap_index = (hash[3] >> 16) & BITMAP_MASK;
	tmp &= (bitmaps[BITMAP_SHIFT * 7 + (bitmap_index >> 5)] >>
	        (bitmap_index & 31)) & 1U;
#elif SELECT_CMP_STEPS > 2
	bitmap_index = hash[3] & BITMAP_MASK;
	tmp &= (bitmaps[bitmap_index >> 5] >> (bitmap_index & 31)) & 1U;
	bitmap_index = hash[2] & BITMAP_MASK;
	tmp &= (bitmaps[BITMAP_SHIFT + (bitmap_index >> 5)] >> (bitmap_index & 31)) &
	       1U;
	bitmap_index = hash[1] & BITMAP_MASK;
	tmp &= (bitmaps[BITMAP_SHIFT * 2 + (bitmap_index >> 5)] >>
	        (bitmap_index & 31)) & 1U;
	bitmap_index = hash[0] & BITMAP_MASK;
	tmp &= (bitmaps[BITMAP_SHIFT * 3 + (bitmap_index >> 5)] >>
	        (bitmap_index & 31)) & 1U;
#elif SELECT_CMP_STEPS > 1
	bitmap_index = hash[3] & BITMAP_MASK;
	tmp &= (bitmaps[bitmap_index >> 5] >> (bitmap_index & 31)) & 1U;
	bitmap_index = hash[2] & BITMAP_MASK;
	tmp &= (bitmaps[BITMAP_SHIFT + (bitmap_index >> 5)] >> (bitmap_index & 31)) &
	       1U;
#else
	bitmap_index = hash[3] & BITMAP_MASK;
	tmp &= (bitmaps[bitmap_index >> 5] >> (bitmap_index & 31)) & 1U;
#endif

	if (tmp)
		cmp_final(gid, iter, hash, offset_table, hash_table, return_hashes, output,
		          bitmap_dupe);
}

#define USE_CONST_CACHE \
    (CONST_CACHE_SIZE >= (NUM_INT_KEYS * 4))

typedef struct {
	uint iter;
	uint len;
	uchar salt[MAX_SALT_SIZE];
} salt_t;

__kernel
void sha256(__global uint *keys,
            __global uint *index,
            __global uint *int_key_loc,
#if USE_CONST_CACHE
            constant
#else
            __global
#endif
            uint *int_keys,
            __constant salt_t *salt,
            __global uint *bitmaps,
            __global uint *offset_table,
            __global uint *hash_table,
            __global uint *return_hashes,
            volatile __global uint *out_hash_ids,
            volatile __global uint *bitmap_dupe)
{
	uint i;
	uint gid = get_global_id(0);
	uint base = index[gid];
	uint len = base & 63;
	uint T[16] = { 0 };

#if NUM_INT_KEYS > 1 && !IS_STATIC_GPU_MASK
	uint ikl = int_key_loc[gid];
	uint loc0 = ikl & 0xff;
#if MASK_FMT_INT_PLHDR > 1
#if LOC_1 >= 0
	uint loc1 = (ikl & 0xff00) >> 8;
#endif
#endif
#if MASK_FMT_INT_PLHDR > 2
#if LOC_2 >= 0
	uint loc2 = (ikl & 0xff0000) >> 16;
#endif
#endif
#if MASK_FMT_INT_PLHDR > 3
#if LOC_3 >= 0
	uint loc3 = (ikl & 0xff000000) >> 24;
#endif
#endif
#endif

#if !IS_STATIC_GPU_MASK
#define GPU_LOC_0 loc0
#define GPU_LOC_1 loc1
#define GPU_LOC_2 loc2
#define GPU_LOC_3 loc3
#else
#define GPU_LOC_0 LOC_0
#define GPU_LOC_1 LOC_1
#define GPU_LOC_2 LOC_2
#define GPU_LOC_3 LOC_3
#endif

#if USE_LOCAL_BITMAPS
	uint lid = get_local_id(0);
	uint lws = get_local_size(0);
	__local uint s_bitmaps[BITMAP_SHIFT * SELECT_CMP_STEPS];

	for (i = lid; i < BITMAP_SHIFT * SELECT_CMP_STEPS; i += lws)
		s_bitmaps[i] = bitmaps[i];

	barrier(CLK_LOCAL_MEM_FENCE);
#endif

	keys += base >> 6;

	for (i = 0; i < salt->len; i++)
		PUTCHAR_BE(T, i, salt->salt[i]);

	__global uchar *key = (__global uchar *)keys;
	for (; i < salt->len + len; i++)
		PUTCHAR_BE(T, i, *key++);

	PUTCHAR_BE(T, (salt->len + len), 0x80);
	T[15] = (salt->len + len) << 3;

	for (uint idx = 0; idx < NUM_INT_KEYS; idx++) {
#if NUM_INT_KEYS > 1
		PUTCHAR_BE(T, salt->len + GPU_LOC_0, (int_keys[idx] & 0xff));
#if MASK_FMT_INT_PLHDR > 1
#if LOC_1 >= 0
		PUTCHAR_BE(T, salt->len + GPU_LOC_1, ((int_keys[idx] & 0xff00) >> 8));
#endif
#endif
#if MASK_FMT_INT_PLHDR > 2
#if LOC_2 >= 0
		PUTCHAR_BE(T, salt->len + GPU_LOC_2, ((int_keys[idx] & 0xff0000) >> 16));
#endif
#endif
#if MASK_FMT_INT_PLHDR > 3
#if LOC_3 >= 0
		PUTCHAR_BE(T, salt->len + GPU_LOC_3, ((int_keys[idx] & 0xff000000) >> 24));
#endif
#endif
#endif
		uint W[16];
		uint hash[8];

		memcpy_macro(W, T, 16);
		sha256_single(W, hash);

		uint iter = salt->iter;
		while (--iter) {
			// Copy hash directly for next iteration
			memcpy_macro(W, hash, 8);
			W[8] = 0x80000000U;
			W[9] = 0;
			W[10] = 0;
			W[11] = 0;
			W[12] = 0;
			W[13] = 0;
			W[14] = 0;
			W[15] = 32 << 3;
			sha256_single(W, hash);
		}

		// Swap hash to match endianness of bitmap/hash table
		for (uint j = 0; j < 8; j++)
			hash[j] = SWAP32(hash[j]);

		cmp(gid, idx, hash,
#if USE_LOCAL_BITMAPS
		    s_bitmaps
#else
		    bitmaps
#endif
		    , offset_table, hash_table, return_hashes, out_hash_ids, bitmap_dupe);
	}
}
