package crypto

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"os"
)

type KeyPair struct {
	privateKeyBase64 RSAKey
	publicKeyBase64  RSAKey
	private          *rsa.PrivateKey
	public           *rsa.PublicKey
}

func (kp KeyPair) PrivateKeyBase64() RSAKey {
	return kp.privateKeyBase64
}

func (kp KeyPair) PublicKeyBase64() RSAKey {
	return kp.publicKeyBase64
}

func (kp KeyPair) Random() (*KeyPair, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, err
	}

	kp.privateKeyBase64, kp.publicKeyBase64 = KeysToBase64(privateKey)
	kp.private = privateKey
	kp.public = &privateKey.PublicKey

	return &kp, nil
}

func (kp KeyPair) FromBase64(privateKeyBase64 string) (*KeyPair, error) {
	privateKey, err := Base64ToKeys(privateKeyBase64)
	if err != nil {
		return nil, err
	}

	kp.privateKeyBase64, kp.publicKeyBase64 = KeysToBase64(privateKey)
	kp.private = privateKey
	kp.public = &privateKey.PublicKey

	return &kp, nil
}

func (kp *KeyPair) Save(dir string) {
	priv, err := os.Create(dir + "/private.rsa")
	throw(err)
	priv.WriteString(string(kp.privateKeyBase64))
	defer priv.Close()

	pub, err := os.Create(dir + "/public.rsa")
	throw(err)
	pub.WriteString(string(kp.publicKeyBase64))
	defer pub.Close()

	data, err := x509.MarshalPKIXPublicKey(kp.public)

	throw(err)
	pemkey := &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: data,
	}
	pubPem, err := os.Create(dir + "/public.pem")
	throw(err)
	pem.Encode(pubPem, pemkey)

	defer pubPem.Close()
	data, err = x509.MarshalPKCS8PrivateKey(kp.private)
	throw(err)
	pemkey = &pem.Block{
		Type:  "PRIVATE KEY",
		Bytes: data,
	}
	privPem, err := os.Create(dir + "/private.pem")
	throw(err)
	pem.Encode(privPem, pemkey)

	defer privPem.Close()
}

func (kp KeyPair) Load(dir string) *KeyPair {
	priv, err := os.ReadFile(dir + "/private.rsa")
	throw(err)
	data, _ := kp.FromBase64(string(priv))
	return data
}