package encryptor

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"

	"github.com/jchavannes/go-pgp/pgp"
	"golang.org/x/crypto/openpgp"
	"golang.org/x/crypto/openpgp/armor"
)

// Encryptor interface
type Encryptor interface {
	Encrypt(string) string
	Decrypt() string
}

type encryptor struct {
	entity *openpgp.Entity
}

type readerWriter struct {
	b []byte
}

func (rw *readerWriter) Write(b []byte) (int, error) {
	rw.b = b[:]
	return len(b), io.EOF
}

func (rw *readerWriter) Read(b []byte) (int, error) {
	b = rw.b[:]
	return len(b), io.EOF
}

func NewKeyPair() (string, string) {
	entity, err := openpgp.NewEntity("a", "b", "c@d.e", nil)
	if err != nil {
		panic(err)
	}

	// Sign all the identities
	for _, id := range entity.Identities {
		err := id.SelfSignature.SignUserId(id.UserId.Id, entity.PrimaryKey, entity.PrivateKey, nil)
		if err != nil {
			panic(err)
		}
	}

	rw := &readerWriter{b: make([]byte, 5000)}
	b := make([]byte, 5000)

	priW, err := armor.Encode(rw, openpgp.PrivateKeyType, nil)
	if err != io.EOF {
		fmt.Println(priW, rw)
		panic(err)
	}
	if priW != nil {
		defer priW.Close()
		entity.SerializePrivate(priW, nil)
	}
	rw.Read(b)
	pri := string(b)

	pubW, err := armor.Encode(rw, openpgp.PublicKeyType, nil)
	if err != io.EOF {
		panic(err)
	}
	rw.Read(b)
	pub := string(b)
	if pubW != nil {
		defer pubW.Close()
		entity.Serialize(pubW)
	}
	fmt.Println(pri, pub)

	return pri, pub
}

// NewEncryptor function
func NewEncryptor(private, public string) *encryptor {
	pubKeyPack, err := pgp.GetPublicKeyPacket([]byte(public))
	if err != nil {
		panic(err)
	}
	priKeyPack, err := pgp.GetPrivateKeyPacket([]byte(private))
	if err != nil {
		panic(err)
	}
	entity, err := openpgp.NewEntity("", "", "", nil)
	entity.PrimaryKey = pubKeyPack
	entity.PrivateKey = priKeyPack
	return &encryptor{
		entity: entity,
	}
}

func b64enc(msg io.Reader) (string, error) {
	bytes, err := ioutil.ReadAll(msg)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(bytes), nil
}

func b64dec(msg string) (io.Reader, error) {
	decMe, err := base64.StdEncoding.DecodeString(msg)
	if err != nil {
		return nil, err
	}
	return bytes.NewBuffer(decMe), nil
}

// Encrypt function
func (e *encryptor) Encrypt(msg string) (string, error) {
	// encrypt string
	buf := new(bytes.Buffer)
	w, err := openpgp.Encrypt(buf, []*openpgp.Entity{e.entity}, nil, nil, nil)
	if err != nil {
		return "", err
	}
	_, err = w.Write([]byte(msg))
	if err != nil {
		return "", err
	}
	err = w.Close()
	if err != nil {
		return "", err
	}
	// Encode to base64
	return b64enc(buf)
}

// Decrypt function
func (e *encryptor) Decrypt(msg string) (string, error) {
	// Decode the base64 string
	decMe, err := b64dec(msg)
	if err != nil {
		panic(err)
	}

	md, err := openpgp.ReadMessage(decMe, openpgp.EntityList([]*openpgp.Entity{e.entity}), nil, nil)
	if err != nil {
		return "", err
	}

	plaintext, err := ioutil.ReadAll(md.UnverifiedBody)
	if err != nil {
		return "", err
	}

	return string(plaintext), nil
}