/**
 * @file SHA.cpp
 * @author apio (cloudapio.eu)
 * @brief SHA256 algorithm implementation.
 *
 * @copyright Copyright (c) 2024, the Luna authors.
 *
 */

#include <endian.h>
#include <luna/Alignment.h>
#include <luna/Common.h>
#include <luna/DebugLog.h>
#include <luna/SHA.h>
#include <luna/StringBuilder.h>
#include <string.h>

// Taken from https://gist.github.com/hak8or/8794351
#define ROTRIGHT(word, bits) (((word) >> (bits)) | ((word) << (32 - (bits))))

u32 k[64] = { 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 };

u32 static H0 = 0x6a09e667;
u32 static H1 = 0xbb67ae85;
u32 static H2 = 0x3c6ef372;
u32 static H3 = 0xa54ff53a;
u32 static H4 = 0x510e527f;
u32 static H5 = 0x9b05688c;
u32 static H6 = 0x1f83d9ab;
u32 static H7 = 0x5be0cd19;

Result<void> SHA256::append(const u8* data, usize size)
{
    return m_message.append_data(data, size);
}

// Thank you to https://sha256algorithm.com/ for showing me how the SHA256 algorithm works, and effectively, how to
// implement it!
Result<SHA256::Digest> SHA256::digest()
{
    usize original_size = m_message.size();
    usize message_block_length = original_size + 1 + sizeof(u64);
    usize message_block_chunks = ceil_div(message_block_length, 64ul);
    usize num_zeros = message_block_chunks * 64 - message_block_length;

    // Add one bit at the end.
    u8 one = 0b10000000;
    TRY(m_message.append_data(&one, 1));

    // Fill with zeros.
    u8* slice = TRY(m_message.slice_at_end(num_zeros));
    memset(slice, 0, num_zeros);

    // Add the length of the original message (in bits), this has to be big-endian.
    usize message_length = original_size * 8;
    message_length = htobe64(message_length);
    TRY(m_message.append_data((const u8*)&message_length, sizeof(usize)));

    // The length of the message block should now be a multiple of 512 bits (64 bytes).
    check(is_aligned<64>(m_message.size()));

    u32 a, b, c, d, e, f, g, h;         // working variables
    u32 h0, h1, h2, h3, h4, h5, h6, h7; // hash values

    // Initialize the hash values to their starting values.
    h0 = H0;
    h1 = H1;
    h2 = H2;
    h3 = H3;
    h4 = H4;
    h5 = H5;
    h6 = H6;
    h7 = H7;

    // For each 512-bit (64-byte) chunk, perform the following operations.
    for (usize i = 0; i < message_block_chunks; i++)
    {
        // Create a message schedule of 64 dwords, and copy the current chunk into the first 16 dwords.
        u32 message_schedule[64];

        memcpy(message_schedule, &m_message.data()[i * 64], 64);

#if BYTE_ORDER == LITTLE_ENDIAN
        for (int j = 0; j < 16; j++)
        {
            // SHA uses big-endian.
            message_schedule[j] = htobe32(message_schedule[j]);
        }
#endif

        // Fill the rest of the message schedule.
        for (int j = 0; j < 48; j++)
        {
            u32 w0 = message_schedule[j];
            u32 w1 = message_schedule[j + 1];
            u32 w14 = message_schedule[j + 14];
            u32 w9 = message_schedule[j + 9];

            u32 s0 = ROTRIGHT(w1, 7) ^ ROTRIGHT(w1, 18) ^ (w1 >> 3);
            u32 s1 = ROTRIGHT(w14, 17) ^ ROTRIGHT(w14, 19) ^ (w14 >> 10);

            message_schedule[j + 16] = w0 + s0 + w9 + s1;
        }

        // Initialize the working variables with the current hash values.
        a = h0;
        b = h1;
        c = h2;
        d = h3;
        e = h4;
        f = h5;
        g = h6;
        h = h7;

        // For each dword in the message schedule, update the working variables, using this algorithm and the
        // corresponding dword from the constant k[] array.
        for (int j = 0; j < 64; j++)
        {
            u32 temp1, temp2, choice, majority;
            choice = (e & f) ^ ((~e) & g);
            majority = (a & b) ^ (a & c) ^ (b & c);
            temp1 = h + (ROTRIGHT(e, 6) ^ ROTRIGHT(e, 11) ^ ROTRIGHT(e, 25)) + choice + k[j] + message_schedule[j];
            temp2 = (ROTRIGHT(a, 2) ^ ROTRIGHT(a, 13) ^ ROTRIGHT(a, 22)) + majority;

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

        // Add the working variables to the hash values.
        h0 += a;
        h1 += b;
        h2 += c;
        h3 += d;
        h4 += e;
        h5 += f;
        h6 += g;
        h7 += h;
    }

    h0 = be32toh(h0);
    h1 = be32toh(h1);
    h2 = be32toh(h2);
    h3 = be32toh(h3);
    h4 = be32toh(h4);
    h5 = be32toh(h5);
    h6 = be32toh(h6);
    h7 = be32toh(h7);

    Digest result;
    memcpy(&result.m_bytes[0], &h0, sizeof(u32));
    memcpy(&result.m_bytes[4], &h1, sizeof(u32));
    memcpy(&result.m_bytes[8], &h2, sizeof(u32));
    memcpy(&result.m_bytes[12], &h3, sizeof(u32));
    memcpy(&result.m_bytes[16], &h4, sizeof(u32));
    memcpy(&result.m_bytes[20], &h5, sizeof(u32));
    memcpy(&result.m_bytes[24], &h6, sizeof(u32));
    memcpy(&result.m_bytes[28], &h7, sizeof(u32));

    m_message.try_resize(0); // Don't need this anymore.

    return result;
}

Result<String> SHA256::Digest::to_string()
{
    StringBuilder sb;

    for (usize i = 0; i < sizeof(m_bytes); i++) { sb.format("%02x", m_bytes[i]); }

    return sb.string();
}