module angel.utils.cryptography.cryptography;

// Internal imports
import angel.utils.logging;
import angel.utils.cryptography.serpent;
import angel.utils.cryptography.threefish;
import angel.utils.cryptography.curve25519;
// External imports
import std.stdio;
import std.random;
import std.format;

class Cryptography {
    public {
        struct KeyPair {
            ubyte[32] clientSecretKey;
            ubyte[32] sharedSecret;
        }
    }

    public static KeyPair derive_25519(ubyte[] pk) {
        ubyte[32] sk; // generate client secret key
        for (int i = 0; i < 32; ++i) {
            sk[i] = cast(ubyte)(uniform(0, 256));
        }

        Logger.log(LogLevel.Debug, "Generated client sk");

        ubyte[32] ss = curve25519_scalarmult(sk, pk); // derive shared secret out of pk and sk

        Logger.log(LogLevel.Debug, format("Derived shared secret: %s", ss));

        return KeyPair(sk, ss);
    }


}

class Freefishy {
    private Threefish512 cipher;

    this(ulong[8] key, ulong[2] tweak) {
        cipher = new Threefish512();
        cipher.setup(key, tweak);
    }

    public ubyte[] encrypt(ubyte[] payload) {
        ulong[8] plain = [0, 0, 0, 0, 0, 0, 0, 0];

        // Pack the ubyte[] into ulong[] by 8 bytes per ulong
        foreach (i, c; payload) {
            plain[i / 8] |= cast(ulong) c << ((i % 8) * 8);
        }

        // Encrypt the ulong array
        ulong[] encrypted_three = cipher.crypt(plain);

        // Convert the encrypted ulong[] back to ubyte[]
        ubyte[] encrypted_bytes;
        foreach (enc_val; encrypted_three) {
            foreach (i; 0..8) {
                encrypted_bytes ~= cast(ubyte)((enc_val >> (i * 8)) & 0xFF);
            }
        }

        Logger.log(LogLevel.Debug, format("Threefish Encrypted: %s", encrypted_bytes));
        return encrypted_bytes;
    }

    public ubyte[] decrypt(ubyte[] payload) {
        ulong[8] decrypted_buffer;

        // Convert the ubyte[] payload back into ulong[] (8 bytes per ulong)
        ulong[8] cipher_payload;
        foreach (i, val; payload) {
            cipher_payload[i / 8] |= cast(ulong) val << ((i % 8) * 8);
        }

        // Decrypt the ulong[] array
        cipher.decrypt(cipher_payload.ptr, decrypted_buffer.ptr);

        // Convert the decrypted ulong[] back to ubyte[]
        ubyte[] decrypted_bytes;
        foreach (dec_val; decrypted_buffer) {
            foreach (i; 0..8) {
                decrypted_bytes ~= cast(ubyte)((dec_val >> (i * 8)) & 0xFF);
            }
        }

        Logger.log(LogLevel.Debug, format("Threefish Decrypted: %s", decrypted_bytes));
        return decrypted_bytes;
    }
}


class Serpi {
    private Serpent serp;

    this(Cryptography.KeyPair keypair) {
        auto key = cast(ubyte[]) keypair.sharedSecret.dup;
        serp.start(key);
    }

    public ubyte[] encrypt(ubyte[] payload) {
        ubyte padding = cast(ubyte)(16 - (payload.length % 16));
        ubyte[] paddingBytes = new ubyte[padding];
        foreach (i; 0 .. padding) {
            paddingBytes[i] = padding;
        }
        ubyte[] input = payload ~ paddingBytes;
        ubyte[] output = new ubyte[input.length];

        serp.encrypt(input, output);
        Logger.log(LogLevel.Debug, format("Serpent Encrypted data: %s", output));

        return output;
    }

    public ubyte[] decrypt(ubyte[] payload) {
        ubyte[] output = new ubyte[payload.length];
        serp.decrypt(payload, output);

        if (output.length == 0) {
            throw new Exception("Decryption failed: Output length is zero.");
        }

        ubyte padding = output[$ - 1];

        if (padding > 16 || padding == 0) {
            throw new Exception("Invalid padding detected");
        }

        foreach (i; output.length - padding .. output.length) {
            if (output[i] != padding) {
                throw new Exception("Padding validation failed");
            }
        }
        
        ubyte[] unpadded = output[0 .. output.length - padding];
        Logger.log(LogLevel.Debug, format("Serpent Decrypted data: %s", unpadded));

        return unpadded;
    }

    public void reset() {
        serp.reset();
    }
}