package ganarchy.friendcode.util;

import org.apache.commons.codec.binary.Base64;

import javax.crypto.*;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.*;
import java.util.Properties;

/**
 * Helper to deal with private keys.
 */
public class KeyUtil {
    // it is important to use encryption, as you don't want random users
    // hosting the world code. by encrypting the world code to an user, it
    // prevents that.
    // then you can freely share the world files and nothing bad happens.

    /**
     * Reads and decrypts a private key file.
     *
     * @param keyFile The path to the key file.
     * @return The decrypted private key, or null if it doesn't exist.
     */
    public static String readKeyFile(Path keyFile) {
        byte[] file;
        byte[] out;
        try {
            file = Files.readAllBytes(keyFile);
        } catch (IOException e) {
            return null;
        }
        if (file.length < 12 + 16 + 884) { // 12 (IV) + 16 (tag) + 884 (key)
            return null;
        }
        try {
            var c = Cipher.getInstance("AES/GCM/NoPadding");
            Key key = getGlobalKey();
            byte[] iv = new byte[12];
            System.arraycopy(file, 0, iv, 0, 12);
            c.init(Cipher.DECRYPT_MODE, key, new GCMParameterSpec(128, iv));
            out = c.doFinal(file, 12, file.length - 12);
        } catch (AEADBadTagException e) {
            return null;
        } catch (
            NoSuchAlgorithmException | NoSuchPaddingException
            | InvalidAlgorithmParameterException | InvalidKeyException
            | IllegalBlockSizeException | BadPaddingException e
        ) {
            throw new RuntimeException(e);
        }
        return StandardCharsets.UTF_8.decode(ByteBuffer.wrap(out)).toString();
    }


    /**
     * Encrypts and writes a private key file.
     *
     * @param keyFile    The path to the key file.
     * @param privateKey The private key.
     * @return Whether writing the key succeeded.
     */
    public static boolean writeKeyFile(Path keyFile, String privateKey) {
        byte[] in = privateKey.getBytes(StandardCharsets.UTF_8);
        byte[] iv = new byte[12];
        byte[] out;
        try {
            var c = Cipher.getInstance("AES/GCM/NoPadding");
            Key key = getGlobalKey();
            var rand = new SecureRandom();
            rand.nextBytes(iv);
            c.init(Cipher.ENCRYPT_MODE, key, new GCMParameterSpec(128, iv));
            out = c.doFinal(in, 0, in.length);
        } catch (
            NoSuchAlgorithmException | NoSuchPaddingException
            | InvalidAlgorithmParameterException | InvalidKeyException
            | IllegalBlockSizeException | BadPaddingException e
        ) {
            throw new RuntimeException(e);
        }
        try (var file = Files.newOutputStream(keyFile)) {
            file.write(iv);
            file.write(out);
            file.flush();
            return true;
        } catch (IOException e) {
            return false;
        }
    }

    /**
     * Returns the global key used for the encryption/decryption of world keys.
     */
    private static Key getGlobalKey() {
        Properties prop = new Properties();
        if (ConfigUtil.getSettings(prop)) {
            var key = prop.getProperty("key.aes");
            if (key != null) {
                byte[] bytes = null;
                try {
                    bytes = Base64.decodeBase64(key);
                } catch (IllegalArgumentException ignored) {
                }
                if (bytes != null && bytes.length == 16) {
                    return new SecretKeySpec(bytes, "AES");
                }
            }
        }
        prop.clear();
        try {
            var gen = KeyGenerator.getInstance("AES");
            gen.init(128);
            Key key = gen.generateKey();
            var encoded = Base64.encodeBase64String(key.getEncoded());
            prop.setProperty("key.aes", encoded);
            // if the key didn't save we'll just regen world codes.
            ConfigUtil.updateSettings(prop);
            return key;
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
    }
}