package encrypt

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"os"

	"github.com/OneOfOne/xxhash"
)

type Encrypt struct {
	PrivateKey *rsa.PrivateKey
}

func LoadKey(path string, keygen bool) (encrypt *Encrypt, err error) {
	var privateKey *rsa.PrivateKey
	if keygen {
		privateKey, err = generateKey()
		if err != nil {
			return nil, err
		}
		if err := writeKeyToFile(path, privateKey); err != nil {
			return nil, err
		}
	} else {
		privateKey, err = readPrivateKey(path)
		if err != nil {
			return nil, err
		}
	}
	return &Encrypt{
		PrivateKey: privateKey,
	}, nil
}

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

func writeKeyToFile(path string, privateKey *rsa.PrivateKey) error {
	pemdata := pem.EncodeToMemory(
		&pem.Block{
			Type:  "RSA PRIVATE KEY",
			Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
		},
	)
	if _, err := os.Stat(path); os.IsNotExist(err) {
		if err := ioutil.WriteFile(path, pemdata, 0640); err != nil {
			return err
		}
		return nil
	}
	return fmt.Errorf("private key file %s already exists", path)
}

func readPrivateKey(path string) (*rsa.PrivateKey, error) {
	file, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	pemdata, _ := pem.Decode(file)
	privateKey, err := x509.ParsePKCS1PrivateKey(pemdata.Bytes)
	if err != nil {
		return nil, err
	}
	return privateKey, nil
}

func (enc *Encrypt) Encrypt(secretMessage string) (string, error) {
	ciphertext, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, &enc.PrivateKey.PublicKey, []byte(secretMessage), []byte("OAEP Encrypted"))
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

func (enc *Encrypt) Decrypt(cipherText string) (string, error) {
	ct, err := base64.StdEncoding.DecodeString(cipherText)
	if err != nil {
		return "", err
	}
	plaintext, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, enc.PrivateKey, ct, []byte("OAEP Encrypted"))
	if err != nil {
		return "", err
	}
	return string(plaintext), nil
}

func Hash(key []byte) uint64 {
	hash := xxhash.New64()
	hash.Write(key)
	return hash.Sum64()
}