local Bit = require(".crypto.util.bit")
local Queue = require(".crypto.util.queue")
local Stream = require(".crypto.util.stream")

local CONSTANTS = {
	0x428a2f98,
	0x71374491,
	0xb5c0fbcf,
	0xe9b5dba5,
	0x3956c25b,
	0x59f111f1,
	0x923f82a4,
	0xab1c5ed5,
	0xd807aa98,
	0x12835b01,
	0x243185be,
	0x550c7dc3,
	0x72be5d74,
	0x80deb1fe,
	0x9bdc06a7,
	0xc19bf174,
	0xe49b69c1,
	0xefbe4786,
	0x0fc19dc6,
	0x240ca1cc,
	0x2de92c6f,
	0x4a7484aa,
	0x5cb0a9dc,
	0x76f988da,
	0x983e5152,
	0xa831c66d,
	0xb00327c8,
	0xbf597fc7,
	0xc6e00bf3,
	0xd5a79147,
	0x06ca6351,
	0x14292967,
	0x27b70a85,
	0x2e1b2138,
	0x4d2c6dfc,
	0x53380d13,
	0x650a7354,
	0x766a0abb,
	0x81c2c92e,
	0x92722c85,
	0xa2bfe8a1,
	0xa81a664b,
	0xc24b8b70,
	0xc76c51a3,
	0xd192e819,
	0xd6990624,
	0xf40e3585,
	0x106aa070,
	0x19a4c116,
	0x1e376c08,
	0x2748774c,
	0x34b0bcb5,
	0x391c0cb3,
	0x4ed8aa4a,
	0x5b9cca4f,
	0x682e6ff3,
	0x748f82ee,
	0x78a5636f,
	0x84c87814,
	0x8cc70208,
	0x90befffa,
	0xa4506ceb,
	0xbef9a3f7,
	0xc67178f2,
}

local fmt = "%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x"
	.. "%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x"

local AND = Bit.band
local OR = Bit.bor
local NOT = Bit.bnot
local XOR = Bit.bxor
local RROT = Bit.rrotate
local LSHIFT = Bit.lshift
local RSHIFT = Bit.rshift

--SHA2 is big-endian
local bytes2word = function(b0, b1, b2, b3)
	local i = b0
	i = LSHIFT(i, 8)
	i = OR(i, b1)
	i = LSHIFT(i, 8)
	i = OR(i, b2)
	i = LSHIFT(i, 8)
	i = OR(i, b3)
	return i
end

local word2bytes = function(word)
	local b0, b1, b2, b3
	b3 = AND(word, 0xFF)
	word = RSHIFT(word, 8)
	b2 = AND(word, 0xFF)
	word = RSHIFT(word, 8)
	b1 = AND(word, 0xFF)
	word = RSHIFT(word, 8)
	b0 = AND(word, 0xFF)
	return b0, b1, b2, b3
end

local dword2bytes = function(i)
	local b4, b5, b6, b7 = word2bytes(i)
	local b0, b1, b2, b3 = word2bytes(math.floor(i / 0x100000000))
	return b0, b1, b2, b3, b4, b5, b6, b7
end

local SHA2_256 = function()
	local queue = Queue()

	local h0 = 0x6a09e667
	local h1 = 0xbb67ae85
	local h2 = 0x3c6ef372
	local h3 = 0xa54ff53a
	local h4 = 0x510e527f
	local h5 = 0x9b05688c
	local h6 = 0x1f83d9ab
	local h7 = 0x5be0cd19

	local public = {}

	local processBlock = function()
		local a = h0
		local b = h1
		local c = h2
		local d = h3
		local e = h4
		local f = h5
		local g = h6
		local h = h7

		local w = {}

		for i = 0, 15 do
			w[i] = bytes2word(queue.pop(), queue.pop(), queue.pop(), queue.pop())
		end

		for i = 16, 63 do
			local s0 = XOR(RROT(w[i - 15], 7), XOR(RROT(w[i - 15], 18), RSHIFT(w[i - 15], 3)))
			local s1 = XOR(RROT(w[i - 2], 17), XOR(RROT(w[i - 2], 19), RSHIFT(w[i - 2], 10)))
			w[i] = AND(w[i - 16] + s0 + w[i - 7] + s1, 0xFFFFFFFF)
		end

		for i = 0, 63 do
			local s1 = XOR(RROT(e, 6), XOR(RROT(e, 11), RROT(e, 25)))
			local ch = XOR(AND(e, f), AND(NOT(e), g))
			local temp1 = h + s1 + ch + CONSTANTS[i + 1] + w[i]
			local s0 = XOR(RROT(a, 2), XOR(RROT(a, 13), RROT(a, 22)))
			local maj = XOR(AND(a, b), XOR(AND(a, c), AND(b, c)))
			local temp2 = s0 + maj

			h = g
			g = f
			f = e
			e = d + temp1
			d = c
			c = b
			b = a
			a = temp1 + temp2
		end

		h0 = AND(h0 + a, 0xFFFFFFFF)
		h1 = AND(h1 + b, 0xFFFFFFFF)
		h2 = AND(h2 + c, 0xFFFFFFFF)
		h3 = AND(h3 + d, 0xFFFFFFFF)
		h4 = AND(h4 + e, 0xFFFFFFFF)
		h5 = AND(h5 + f, 0xFFFFFFFF)
		h6 = AND(h6 + g, 0xFFFFFFFF)
		h7 = AND(h7 + h, 0xFFFFFFFF)
	end

	public.init = function()
		queue.reset()

		h0 = 0x6a09e667
		h1 = 0xbb67ae85
		h2 = 0x3c6ef372
		h3 = 0xa54ff53a
		h4 = 0x510e527f
		h5 = 0x9b05688c
		h6 = 0x1f83d9ab
		h7 = 0x5be0cd19

		return public
	end

	public.update = function(bytes)
		for b in bytes do
			queue.push(b)
			if queue.size() >= 64 then
				processBlock()
			end
		end

		return public
	end

	public.finish = function()
		local bits = queue.getHead() * 8

		queue.push(0x80)
		while ((queue.size() + 7) % 64) < 63 do
			queue.push(0x00)
		end

		local b0, b1, b2, b3, b4, b5, b6, b7 = dword2bytes(bits)

		queue.push(b0)
		queue.push(b1)
		queue.push(b2)
		queue.push(b3)
		queue.push(b4)
		queue.push(b5)
		queue.push(b6)
		queue.push(b7)

		while queue.size() > 0 do
			processBlock()
		end

		return public
	end

	public.asBytes = function()
		local b0, b1, b2, b3 = word2bytes(h0)
		local b4, b5, b6, b7 = word2bytes(h1)
		local b8, b9, b10, b11 = word2bytes(h2)
		local b12, b13, b14, b15 = word2bytes(h3)
		local b16, b17, b18, b19 = word2bytes(h4)
		local b20, b21, b22, b23 = word2bytes(h5)
		local b24, b25, b26, b27 = word2bytes(h6)
		local b28, b29, b30, b31 = word2bytes(h7)

		return {
			b0,
			b1,
			b2,
			b3,
			b4,
			b5,
			b6,
			b7,
			b8,
			b9,
			b10,
			b11,
			b12,
			b13,
			b14,
			b15,
			b16,
			b17,
			b18,
			b19,
			b20,
			b21,
			b22,
			b23,
			b24,
			b25,
			b26,
			b27,
			b28,
			b29,
			b30,
			b31,
		}
	end

	public.asHex = function()
		local b0, b1, b2, b3 = word2bytes(h0)
		local b4, b5, b6, b7 = word2bytes(h1)
		local b8, b9, b10, b11 = word2bytes(h2)
		local b12, b13, b14, b15 = word2bytes(h3)
		local b16, b17, b18, b19 = word2bytes(h4)
		local b20, b21, b22, b23 = word2bytes(h5)
		local b24, b25, b26, b27 = word2bytes(h6)
		local b28, b29, b30, b31 = word2bytes(h7)

		return string.format(
			fmt,
			b0,
			b1,
			b2,
			b3,
			b4,
			b5,
			b6,
			b7,
			b8,
			b9,
			b10,
			b11,
			b12,
			b13,
			b14,
			b15,
			b16,
			b17,
			b18,
			b19,
			b20,
			b21,
			b22,
			b23,
			b24,
			b25,
			b26,
			b27,
			b28,
			b29,
			b30,
			b31
		)
	end

	public.asString = function()
		local b0, b1, b2, b3 = word2bytes(h0)
		local b4, b5, b6, b7 = word2bytes(h1)
		local b8, b9, b10, b11 = word2bytes(h2)
		local b12, b13, b14, b15 = word2bytes(h3)
		local b16, b17, b18, b19 = word2bytes(h4)
		local b20, b21, b22, b23 = word2bytes(h5)
		local b24, b25, b26, b27 = word2bytes(h6)
		local b28, b29, b30, b31 = word2bytes(h7)

		return string.pack(
			string.rep("B", 32),
			b0,
			b1,
			b2,
			b3,
			b4,
			b5,
			b6,
			b7,
			b8,
			b9,
			b10,
			b11,
			b12,
			b13,
			b14,
			b15,
			b16,
			b17,
			b18,
			b19,
			b20,
			b21,
			b22,
			b23,
			b24,
			b25,
			b26,
			b27,
			b28,
			b29,
			b30,
			b31
		)
	end

	return public
end

--- @class Stream : table

--- @param stream (Stream) - A function that returns the next byte of the stream, or nil if the stream has ended.
--- @returns table - A table containing the hash in bytes, string, and hex formats.
local sha2_256 = function(stream)
	local result = SHA2_256().update(stream).finish()
	return result
end

return {
	sha2_256 = sha2_256,
	SHA2_256 = SHA2_256,
}
