module crypto.threefish512_ctr.main;

import std.random : Random, unpredictableSeed, uniform;
import std.algorithm : min;
import std.conv : emplace;
// import dmd.common.blake3;
import std.array : appender;
import std.exception : enforce;

/* honestly dk if this is cryptographically "secure", as i took most of the implmentation from some russian dub package
i then implemented pkcs5 back then as thats all i knew
then figured just that is not enough and i only had raw stream ciphers so far
so i tried to research how to implement ctr and shit was too complicated so i vibecoded it (probably some flaws and inefficiet, but should be good enough)
and added blake3-mac as without it the implementation would be vulnerable to
replay, padding‑free oracles, bitflip and plaintext-recovery attacks */

// anyways, from what i read forging a blake3 mac or breaking threefish512 in itself
// is beyond fantasy rn so anyone who wants to use this should be fine

// memcpy
extern(C) nothrow @nogc void* memcpy(void* dst, const void* src, size_t n);

extern(C)
{
    struct blake3_hasher {
        ubyte[256] opaque;
    }

    void blake3_hasher_init_keyed(blake3_hasher* hasher, const ubyte* key);
    void blake3_hasher_update(blake3_hasher* hasher, const ubyte* input, size_t len);
    void blake3_hasher_finalize(blake3_hasher* hasher, ubyte* _out, size_t out_len);
}

class Threefish512
{
    private:
    enum BLOCK_SIZE = 64;
    enum Nw = 8;
    enum Nr = 72;
    enum Ns = Nr / 4;

    uint[8] p   = [2, 1, 4, 7, 6, 5, 0, 3];
    uint[8] p_1 = [6, 1, 0, 7, 2, 5, 4, 3];

    uint[4][8] r = [
        [46, 36, 19, 37], 
        [33, 27, 14, 42], 
        [17, 49, 36, 39], 
        [44, 9 , 54, 56],
        [39, 30, 34, 24], 
        [13, 50, 10, 17], 
        [25, 29, 39, 43],
        [8 , 35, 56, 22]
    ];

    ulong[3] t;
    ulong[8][Ns + 1] subKeys;

    alias _mod8 = (ulong a) => a & 7UL;

    private void _mix(ref ulong[2] x, ulong r, ref ulong[2] y)
    {
        y[0] = x[0] + x[1];
        y[1] = (x[1] << r) | (x[1] >> (64 - r));
        y[1] ^= y[0];
    }

    private void _demix(ref ulong[2] y, ulong r, ref ulong[2] x)
    {
        y[1] ^= y[0];
        x[1] = (y[1] << (64 - r)) | (y[1] >> r);
        x[0] = y[0] - x[1];
    }

    private void _setup(ulong* keyData, ulong* tweakData) @system
    {
        ulong[8] K;
        ulong[2] T;
        ulong[9] key;
        ulong kNw = 0x1BD11BDAA9FC1A22;

        memcpy(&K[0], keyData, 64);
        memcpy(&T[0], tweakData, 16);

        foreach (i; 0 .. Nw)
        {
            kNw ^= K[i];
            key[i] = K[i];
        }
        key[8] = kNw;

        t[0] = T[0];
        t[1] = T[1];
        t[2] = T[0] ^ T[1];

        foreach (round; 0 .. Ns + 1)
        {
            foreach (i; 0 .. Nw)
            {
                subKeys[round][i] = key[(round + i) % (Nw + 1)];
                if (i == Nw - 3) subKeys[round][i] += t[round % 3];
                else if (i == Nw - 2) subKeys[round][i] += t[(round + 1) % 3];
                else if (i == Nw - 1) subKeys[round][i] += round;
            }
        }
    }

    private void _blockEncrypt(ulong* plainData, ulong* c)
    {
        ulong[8] f;
        ulong[8] e;
        ulong[2] x, y;
        ulong[8] v;
        memcpy(&v[0], plainData, 64);

        foreach (round; 0 .. Nr)
        {
            uint s = round >> 2;
            foreach (i; 0 .. Nw)
                e[i] = (round % 4 == 0) ? v[i] + subKeys[s][i] : v[i];

            foreach (i; 0 .. Nw/2)
            {
                x[0] = e[i*2]; x[1] = e[i*2+1];
                _mix(x, r[_mod8(round)][i], y);
                f[i*2] = y[0]; f[i*2+1] = y[1];
            }

            foreach (i; 0 .. Nw) v[i] = f[p[i]];
        }

        foreach (i; 0 .. Nw) c[i] = v[i] + subKeys[Ns][i];
    }

    public static ulong[8] generateKey()
    {
        ulong[8] key;
        auto rng = Random(unpredictableSeed);
        foreach (i; 0 .. 8) key[i] = uniform!ulong(rng);
        return key;
    }

    public static ulong generateNonce()
    {
        auto rng = Random(unpredictableSeed);
        return uniform!ulong(rng);
    }

	public static ubyte[32] generateMacKey() {
		auto rng = Random(unpredictableSeed);
		ubyte[32] key;
		foreach(i; 0 .. 4) {
			ulong part = uniform!ulong(rng);
			foreach(b; 0 .. 8)
				key[i*8 + b] = cast(ubyte)((part >> (8*b)) & 0xFF);
		}
		return key;
	}

    public void encryptCTR(ulong* plainData, ulong* ciphertext, size_t lengthBytes, ulong[8] key, ulong nonce)
    {
        ulong[2] tweak = [nonce, 0];
        //_setup(key, [nonce, 0]);
        _setup(key.ptr, tweak.ptr);
        ulong[8] zeroBlock = 0;
        size_t numBlocks = (lengthBytes + 63) / 64;
        ulong counter = 0;

        foreach (blockIndex; 0 .. numBlocks)
        {
            t[0] = nonce;
            t[1] = counter;
            t[2] = t[0] ^ t[1];

            ulong[8] keystream;
            _blockEncrypt(zeroBlock.ptr, keystream.ptr);

            size_t start = blockIndex * 8;
            size_t blockULongs = ((lengthBytes - start + 7)/8).min(8);

            foreach (i; 0 .. blockULongs)
                ciphertext[start + i] = plainData[start + i] ^ keystream[i];

            counter++;
        }
    }

	private bool ctEqual(ubyte[32] a, ubyte[32] b)
	{
		ubyte diff = 0;
		foreach(i; 0..32)
			diff |= a[i] ^ b[i];
		return diff == 0;
	}

	public ubyte[32] gen_tag(ubyte[32] key_mac, ubyte[] nonce, ubyte[] ciphertext)
	{
		auto data = appender!(ubyte[])();
		data.put(nonce);
		data.put(ciphertext);

		ubyte[32] comp_tag = blake3_mac(key_mac, data.data);
		
		return comp_tag;
	}

	private bool verify_tag(ubyte[32] key_mac, ubyte[] nonce, ubyte[] ciphertext, ubyte[32] exp_tag)
	{
		ubyte[32] comp_tag = gen_tag(key_mac, nonce, ciphertext);
		
		return ctEqual(comp_tag, exp_tag);
	}

	private ubyte[32] blake3_mac(ubyte[32] key_mac, ubyte[] data)
	{
		blake3_hasher hasher;
		blake3_hasher_init_keyed(&hasher, key_mac.ptr);
		blake3_hasher_update(&hasher, data.ptr, data.length);
		
		ubyte[32] oout;
		blake3_hasher_finalize(&hasher, oout.ptr, 32);
		return oout;
	}

    public void decryptCTR(ulong* ciphertext, ulong* plainData, size_t lengthBytes, ulong[8] key, ulong nonce, ubyte[32] key_mac, ubyte[32] exp_tag)
    {
        ubyte[8] nbytes;
        foreach(i; 0 .. 8) nbytes[i] = cast(ubyte)((nonce >> (8*i)) & 0xFF);
    
        ubyte[] ct_bytes = cast(ubyte[])(cast(ubyte*)ciphertext)[0 .. lengthBytes];
		if (verify_tag(key_mac, nbytes[], ct_bytes, exp_tag)) {
			encryptCTR(ciphertext, plainData, lengthBytes, key, nonce);
		}
    }

	public void encryptBytes(ubyte[] plaintext, out ubyte[] ciphertext, out ubyte[32] tag, out ulong nonce, out ulong[8] key_enc, out ubyte[32] key_mac)
    {
        ciphertext.length = plaintext.length;

        size_t nLongs = (plaintext.length + 7)/8;
        ulong[] plainLongs;
        plainLongs.length = nLongs;
        foreach(i; 0 .. nLongs) {
            ulong val = 0;
            foreach(b; 0 .. 8) {
                size_t idx = i*8 + b;
                if(idx < plaintext.length)
                    val |= cast(ulong)plaintext[idx] << (8*b);
            }
            plainLongs[i] = val;
        }

        ulong[] cipherLongs;
        cipherLongs.length = nLongs;

        encryptCTR(plainLongs.ptr, cipherLongs.ptr, plaintext.length, key_enc, nonce);

        foreach(i; 0 .. nLongs) {
            foreach(b; 0 .. 8) {
                size_t idx = i*8 + b;
                if(idx < plaintext.length)
                    ciphertext[idx] = cast(ubyte)((cipherLongs[i] >> (8*b)) & 0xFF);
            }
        }

        auto data = appender!(ubyte[])();
        foreach(b; 0 .. 8)
            data.put(cast(ubyte)((nonce >> (8*b)) & 0xFF));
        data.put(ciphertext);
        tag = blake3_mac(key_mac, data.data);
	}

	public ubyte[] decryptBytes(ubyte[] ciphertext, ubyte[32] expected_tag, ubyte[32] key_mac, ulong[8] key_enc, ulong nonce)
    {
        auto data = appender!(ubyte[])();
        foreach(b; 0 .. 8)
            data.put(cast(ubyte)((nonce >> (8*b)) & 0xFF));
        data.put(ciphertext);

        enforce(ctEqual(blake3_mac(key_mac, data.data), expected_tag), "MAC verification failed");
        size_t nLongs = (ciphertext.length + 7)/8;
        ulong[] cipherLongs;
        cipherLongs.length = nLongs;
        foreach(i; 0 .. nLongs) {
            ulong val = 0;
            foreach(b; 0 .. 8) {
                size_t idx = i*8 + b;
                if(idx < ciphertext.length)
                    val |= cast(ulong)ciphertext[idx] << (8*b);
            }
            cipherLongs[i] = val;
        }

        ulong[] plainLongs;
        plainLongs.length = nLongs;
        encryptCTR(cipherLongs.ptr, plainLongs.ptr, ciphertext.length, key_enc, nonce);

        ubyte[] plaintext;
        plaintext.length = ciphertext.length;
        foreach(i; 0 .. nLongs) {
            foreach(b; 0 .. 8) {
                size_t idx = i*8 + b;
                if(idx < ciphertext.length)
                    plaintext[idx] = cast(ubyte)((plainLongs[i] >> (8*b)) & 0xFF);
            }
        }
        return plaintext;
    }
}

__gshared Threefish512 tfInstance = new Threefish512();

// Key / MAC / Nonce wrappers
extern(C) export void generateKeyC(ulong* c_out) {
    auto key = Threefish512.generateKey(); // returns ulong[8]
    for(size_t i = 0; i < 8; i++)
        c_out[i] = key[i];
}

extern(C) export ulong generateNonceC() {
    return Threefish512.generateNonce();
}

extern(C) export void generateMacKeyC(ubyte* c_out) {
    auto key = Threefish512.generateMacKey(); // returns ubyte[32]
    for(size_t i = 0; i < 32; i++)
        c_out[i] = key[i];
}

// CTR encryption / decryption wrappers
extern(C) export void encryptCTR_C(ulong* plainData, ulong* ciphertext, size_t lengthBytes,
                                   ulong* key_enc, ulong nonce) {
    ulong[8] k;
    for(size_t i = 0; i < 8; i++) k[i] = key_enc[i];
    tfInstance.encryptCTR(plainData, ciphertext, lengthBytes, k, nonce);
}

extern(C) export void decryptCTR_C(ulong* ciphertext, ulong* plainData, size_t lengthBytes,
                                   ulong* key_enc, ulong nonce,
                                   ubyte* key_mac, ubyte* exp_tag) {
    ulong[8] k; ubyte[32] mac; ubyte[32] tag;
    for(size_t i = 0; i < 8; i++) k[i] = key_enc[i];
    for(size_t i = 0; i < 32; i++) { mac[i] = key_mac[i]; tag[i] = exp_tag[i]; }
    tfInstance.decryptCTR(ciphertext, plainData, lengthBytes, k, nonce, mac, tag);
}

// encryptBytes / decryptBytes wrappers
extern(C) export void encryptBytes_C(ubyte* plaintext, size_t plainLen,
                                     ubyte* ciphertext, ubyte* tag,
                                     ulong* nonceOut, ulong* key_encOut, ubyte* key_macOut) {
    ulong[8] key_enc = Threefish512.generateKey();
    ulong nonce = Threefish512.generateNonce();
    ubyte[32] key_mac = Threefish512.generateMacKey();

    ubyte[] pt; pt.length = plainLen;
    for(size_t i = 0; i < plainLen; i++) pt[i] = plaintext[i];

    ubyte[] ct; ct.length = plainLen;
    ubyte[32] tag_arr;
    tfInstance.encryptBytes(pt, ct, tag_arr, nonce, key_enc, key_mac);

    for(size_t i = 0; i < 32; i++) tag[i] = tag_arr[i];

    for(size_t i = 0; i < plainLen; i++) ciphertext[i] = ct[i];
    for(size_t i = 0; i < 8; i++) key_encOut[i] = key_enc[i];
    for(size_t i = 0; i < 32; i++) key_macOut[i] = key_mac[i];
    *nonceOut = nonce;
}

extern(C) export void decryptBytes_C(ubyte* ciphertext, size_t cipherLen,
                                     ubyte* plaintextOut,
                                     ubyte* expectedTag, ubyte* key_mac,
                                     ulong* key_enc, ulong nonce) {
    ulong[8] k; ubyte[32] mac; ubyte[32] tag;
    for(size_t i = 0; i < 8; i++) k[i] = key_enc[i];
    for(size_t i = 0; i < 32; i++) { mac[i] = key_mac[i]; tag[i] = expectedTag[i]; }

    ubyte[] ct; ct.length = cipherLen;
    for(size_t i = 0; i < cipherLen; i++) ct[i] = ciphertext[i];

    ubyte[] pt = tfInstance.decryptBytes(ct, tag, mac, k, nonce);
    for(size_t i = 0; i < cipherLen; i++) plaintextOut[i] = pt[i];
}

// gen_tag wrapper
extern(C) export void gen_tag_C(ubyte* key_mac, ubyte* nonce, size_t nonceLen,
                                ubyte* ciphertext, size_t cipherLen, ubyte* outTag) {
    ubyte[32] k;
    for(size_t i = 0; i < 32; i++) k[i] = key_mac[i];

    ubyte[] n; n.length = nonceLen;
    for(size_t i = 0; i < nonceLen; i++) n[i] = nonce[i];

    ubyte[] ct; ct.length = cipherLen;
    for(size_t i = 0; i < cipherLen; i++) ct[i] = ciphertext[i];

    ubyte[32] tag = tfInstance.gen_tag(k, n, ct);
    for(size_t i = 0; i < 32; i++) outTag[i] = tag[i];
}

extern(C) export void threefish512_block_encrypt_C(
    ulong* key, ulong* tweak, ulong* plaintext, ulong* ciphertext)
{
    ulong[8] k;
    ulong[2] t;

    for(size_t i=0;i<8;i++) k[i] = key[i];
    for(size_t i=0;i<2;i++) t[i] = tweak[i];

    tfInstance._setup(k.ptr, t.ptr);
    tfInstance._blockEncrypt(plaintext, ciphertext);
}