package app.fedilab.android.helper;

import android.content.Context;
import android.content.SharedPreferences;
import android.util.Base64;
import android.util.Log;

import androidx.preference.PreferenceManager;


import org.spongycastle.asn1.ASN1ObjectIdentifier;
import org.spongycastle.asn1.x9.ECNamedCurveTable;
import org.spongycastle.asn1.x9.X9ECParameters;
import org.spongycastle.crypto.params.ECNamedDomainParameters;
import org.spongycastle.crypto.params.ECPublicKeyParameters;
import org.spongycastle.math.ec.ECCurve;
import org.spongycastle.math.ec.ECPoint;

import java.math.BigInteger;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Security;
import java.security.interfaces.ECPublicKey;
import java.security.spec.ECGenParameterSpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.KeyAgreement;

import app.fedilab.android.client.Entities.Account;

// https://github.com/nelenkov/ecdh-kx/blob/master/src/org/nick/ecdhkx/Crypto.java

public class ECDH {


    private static final String kp_public = "kp_public";
    private static final String kp_private = "kp_private";
    private static final String peer_public = "peer_public";

    private static final String TAG = ECDH.class.getSimpleName();

    private static final String PROVIDER = org.spongycastle.jce.provider.BouncyCastleProvider.PROVIDER_NAME;

    private static final String KEGEN_ALG = "ECDH";

    private static ECDH instance;


    static {
        //  Security.removeProvider("BC");
        Security.addProvider(new org.spongycastle.jce.provider.BouncyCastleProvider());
    }

    private final KeyFactory kf;
    private final KeyPairGenerator kpg;


    public ECDH() {
        try {
            kf = KeyFactory.getInstance(KEGEN_ALG, PROVIDER);
            kpg = KeyPairGenerator.getInstance(KEGEN_ALG, PROVIDER);
        } catch (NoSuchAlgorithmException | NoSuchProviderException e) {
            throw new RuntimeException(e);
        }
    }

    public static synchronized ECDH getInstance() {
        if (instance == null) {
            instance = new ECDH();
        }

        return instance;
    }

    public static String base64Encode(byte[] b) {

        byte[] encoded = Base64.encode(
                b, Base64.URL_SAFE | Base64.NO_PADDING | Base64.NO_WRAP);
        return new String(encoded);
    }

    static byte[] base64Decode(String str) {
        return Base64.decode(str, Base64.URL_SAFE | Base64.NO_PADDING | Base64.NO_WRAP);
    }

    synchronized KeyPair generateKeyPair()
            throws Exception {
        ECGenParameterSpec ecParamSpec = new ECGenParameterSpec("prime256v1");
        kpg.initialize(ecParamSpec);

        return kpg.generateKeyPair();
    }

    byte[] generateSecret(PrivateKey myPrivKey, PublicKey otherPubKey) throws Exception {
        ECPublicKey ecPubKey = (ECPublicKey) otherPubKey;
        Log.d(TAG, "public key Wx: "
                + ecPubKey.getW().getAffineX().toString(16));
        Log.d(TAG, "public key Wy: "
                + ecPubKey.getW().getAffineY().toString(16));

        KeyAgreement keyAgreement = KeyAgreement.getInstance(KEGEN_ALG, PROVIDER);
        keyAgreement.init(myPrivKey);
        keyAgreement.doPhase(otherPubKey, true);

        return keyAgreement.generateSecret();
    }

    synchronized PublicKey readPublicKey(String keyStr) throws Exception {
        X509EncodedKeySpec x509ks = new X509EncodedKeySpec(
                base64Decode(keyStr));
        return kf.generatePublic(x509ks);
    }

    synchronized PrivateKey readPrivateKey(String keyStr) throws Exception {
        PKCS8EncodedKeySpec p8ks = new PKCS8EncodedKeySpec(
                base64Decode(keyStr));

        return kf.generatePrivate(p8ks);
    }

    synchronized KeyPair readKeyPair(String pubKeyStr, String privKeyStr)
            throws Exception {
        return new KeyPair(readPublicKey(pubKeyStr), readPrivateKey(privKeyStr));
    }

    KeyPair newPair(Context context) {
        SharedPreferences.Editor prefsEditor = PreferenceManager
                .getDefaultSharedPreferences(context).edit();
        KeyPair kp;
        try {
            kp = generateKeyPair();
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
        ECPublicKey key = (ECPublicKey) kp.getPublic();
        byte[] x = key.getW().getAffineX().toByteArray();
        byte[] y = key.getW().getAffineY().toByteArray();
        BigInteger xbi = new BigInteger(1, x);
        BigInteger ybi = new BigInteger(1, y);
        X9ECParameters x9 = ECNamedCurveTable.getByName("prime256v1");
        ASN1ObjectIdentifier oid = ECNamedCurveTable.getOID("prime256v1");
        ECCurve curve = x9.getCurve();
        ECPoint point = curve.createPoint(xbi, ybi);
        ECNamedDomainParameters dParams = new ECNamedDomainParameters(oid,
                x9.getCurve(), x9.getG(), x9.getN(), x9.getH(), x9.getSeed());
        ECPublicKeyParameters pubKey = new ECPublicKeyParameters(point, dParams);

        String keyString = base64Encode(pubKey.getQ().getEncoded(false));

        prefsEditor.putString(kp_public, keyString);
        prefsEditor.putString(kp_private, base64Encode(kp.getPrivate().getEncoded()));
        prefsEditor.commit();
        return kp;
    }

    synchronized KeyPair getPair(Context context) {
        SharedPreferences prefs = PreferenceManager
                .getDefaultSharedPreferences(context);

        String strPub = prefs.getString(kp_public, "");
        String strPriv = prefs.getString(kp_private, "");

        if (strPub.isEmpty() || strPriv.isEmpty()) {
            return newPair(context);
        }
        try {
            return readKeyPair(strPub, strPriv);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public String getPublicKey(Context context) {
        SharedPreferences prefs = PreferenceManager
                .getDefaultSharedPreferences(context);

        String strPub = prefs.getString(kp_public, "");
        return strPub;
    }

    public void saveServerKey(Context context, Account account, String strPeerPublic) {
        SharedPreferences.Editor prefsEditor = PreferenceManager
                .getDefaultSharedPreferences(context).edit();

        prefsEditor.putString(peer_public + account.getId() + account.getInstance(), strPeerPublic);
        prefsEditor.commit();
    }

    PublicKey getServerKey(Context context, Account account) throws Exception {
        return readPublicKey(
                PreferenceManager.getDefaultSharedPreferences(context)
                        .getString(peer_public + account.getId() + account.getInstance(), "")
        );
    }

    byte[] getSecret(Context context, Account account) {
        try {
            return generateSecret(getPair(context).getPrivate(), getServerKey(context, account));
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }
}