package de.maxhenkel.voicechat.voice.common;

import io.netty.buffer.ByteBuf;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Arrays;
import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

public class Secret {
    public static final int SECRET_SIZE_BYTES = 16;
    public static final int IV_SIZE_BYTES = 12;
    public static final int TAG_LEN_BITS = 128;
    public static final String CIPHER = "AES/GCM/NoPadding";
    private static final SecureRandom RANDOM = new SecureRandom();
    private final byte[] secret;
    private final SecretKeySpec keySpec;
    private volatile boolean legacyProtocol;

    protected Secret(byte[] secret) {
        if (secret == null || secret.length != SECRET_SIZE_BYTES) {
            throw new IllegalArgumentException("Voice chat secret must contain exactly 16 bytes");
        }
        this.secret = secret.clone();
        this.keySpec = new SecretKeySpec(this.secret, "AES");
    }

    public static Secret generateNewRandomSecret() {
        byte[] secret = new byte[SECRET_SIZE_BYTES];
        RANDOM.nextBytes(secret);
        return new Secret(secret);
    }

    public static Secret fromBytes(byte[] secret) {
        return new Secret(secret);
    }

    public static Secret fromBytes(ByteBuf buf) {
        byte[] secretBytes = new byte[SECRET_SIZE_BYTES];
        buf.readBytes(secretBytes);
        return Secret.fromBytes(secretBytes);
    }

    public void toBytes(ByteBuf buf) {
        buf.writeBytes(this.secret);
    }

    public byte[] getSecret() {
        return this.secret.clone();
    }

    public SecretKeySpec getKeySpec() {
        return this.keySpec;
    }

    public void setLegacyProtocol(boolean legacyProtocol) {
        this.legacyProtocol = legacyProtocol;
    }

    public boolean isLegacyProtocol() {
        return this.legacyProtocol;
    }

    public static byte[] generateIV() {
        byte[] iv = new byte[IV_SIZE_BYTES];
        RANDOM.nextBytes(iv);
        return iv;
    }

    private static byte[] generateLegacyIV() {
        byte[] iv = new byte[16];
        RANDOM.nextBytes(iv);
        return iv;
    }

    public byte[] encrypt(byte[] data) throws NoSuchPaddingException, NoSuchAlgorithmException,
            InvalidAlgorithmParameterException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
        if (this.legacyProtocol) {
            byte[] iv = Secret.generateLegacyIV();
            Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
            cipher.init(Cipher.ENCRYPT_MODE, (Key) this.keySpec, new IvParameterSpec(iv));
            byte[] encrypted = cipher.doFinal(data);
            byte[] payload = new byte[iv.length + encrypted.length];
            System.arraycopy(iv, 0, payload, 0, iv.length);
            System.arraycopy(encrypted, 0, payload, iv.length, encrypted.length);
            return payload;
        }
        byte[] iv = Secret.generateIV();
        Cipher cipher = Cipher.getInstance(CIPHER);
        cipher.init(Cipher.ENCRYPT_MODE, (Key) this.keySpec, new GCMParameterSpec(TAG_LEN_BITS, iv));
        byte[] encrypted = cipher.doFinal(data);
        byte[] payload = new byte[iv.length + encrypted.length];
        System.arraycopy(iv, 0, payload, 0, iv.length);
        System.arraycopy(encrypted, 0, payload, iv.length, encrypted.length);
        return payload;
    }

    public byte[] decrypt(byte[] payload) throws NoSuchPaddingException, NoSuchAlgorithmException,
            InvalidAlgorithmParameterException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
        if (this.legacyProtocol) {
            if (payload.length < 32 || (payload.length - 16) % 16 != 0) {
                throw new IllegalBlockSizeException("Invalid legacy AES-CBC packet size");
            }
            byte[] iv = Arrays.copyOfRange(payload, 0, 16);
            byte[] encrypted = Arrays.copyOfRange(payload, 16, payload.length);
            Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
            cipher.init(Cipher.DECRYPT_MODE, (Key) this.keySpec, new IvParameterSpec(iv));
            return cipher.doFinal(encrypted);
        }
        if (payload.length < IV_SIZE_BYTES + 16) {
            throw new IllegalBlockSizeException("Invalid AES-GCM packet size");
        }
        byte[] iv = Arrays.copyOfRange(payload, 0, IV_SIZE_BYTES);
        byte[] encrypted = Arrays.copyOfRange(payload, IV_SIZE_BYTES, payload.length);
        Cipher cipher = Cipher.getInstance(CIPHER);
        cipher.init(Cipher.DECRYPT_MODE, (Key) this.keySpec, new GCMParameterSpec(TAG_LEN_BITS, iv));
        return cipher.doFinal(encrypted);
    }

    @Override
    public boolean equals(Object o) {
        return o instanceof Secret && Arrays.equals(this.secret, ((Secret) o).secret);
    }

    @Override
    public int hashCode() {
        return Arrays.hashCode(this.secret);
    }
}