package encrypter

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"fmt"
	"reflect"
	"strings"

	msgpack "gopkg.in/vmihailenco/msgpack.v2"
)

// PrivateKey type
type PrivateKey rsa.PrivateKey

// PublicKey type
type PublicKey rsa.PublicKey

const (
	keySize          = 2048
	decryptBlockSize = 256
	encryptBlockSize = 190
)

// NewPrivateKey func
func NewPrivateKey() *PrivateKey {
	pk, err := rsa.GenerateKey(rand.Reader, keySize)
	if err != nil {
		panic(fmt.Errorf("Error create rsa.PrivateKey %v", err))
	}
	return (*PrivateKey)(pk)
}

// ParsePrivateKey func
func ParsePrivateKey(raw []byte) *PrivateKey {
	k, err := x509.ParsePKCS1PrivateKey(raw)
	if err != nil {
		panic(fmt.Errorf("Error parse privateKey %v", err))
	}
	return (*PrivateKey)(k)
}

// ParsePrivateKeyString func
func ParsePrivateKeyString(str string) (*PrivateKey, error) {
	raw, err := base64.RawURLEncoding.DecodeString(str)
	if err != nil {
		slog.Debugf("Error key decode %v", err)
		return nil, err
	}
	k, err := x509.ParsePKCS1PrivateKey(raw)
	if err != nil {
		slog.Debugf("Error parse key %v", err)
		return nil, err
	}
	return (*PrivateKey)(k), nil
}

// ParsePublicKey func
func ParsePublicKey(raw []byte) *PublicKey {
	k, err := x509.ParsePKIXPublicKey(raw)
	if err != nil {
		panic(fmt.Errorf("Error parse privateKey %v", err))
	}
	if p, ok := k.(*rsa.PublicKey); ok {
		return (*PublicKey)(p)
	}
	panic(fmt.Errorf("Invalid type %s", reflect.TypeOf(k).String()))
}

// ParsePublicKeyString func
func ParsePublicKeyString(str string) (*PublicKey, error) {
	raw, err := base64.RawURLEncoding.DecodeString(str)
	if err != nil {
		slog.Debugf("Error key decode %v", err)
		return nil, err
	}
	k, err := x509.ParsePKIXPublicKey(raw)
	if err != nil {
		slog.Debugf("Error parse key %v", err)
		return nil, err
	}
	if p, ok := k.(*rsa.PublicKey); ok {
		return (*PublicKey)(p), nil
	}
	err = fmt.Errorf("Invalid type %s", reflect.TypeOf(k).String())
	slog.Debugf("Error %v", err)
	return nil, err
}

var _ msgpack.CustomEncoder = (*PrivateKey)(nil)
var _ msgpack.CustomDecoder = (*PrivateKey)(nil)

// EncodeMsgpack func
func (s *PrivateKey) EncodeMsgpack(enc *msgpack.Encoder) error {
	return enc.Encode(s.Raw())
}

// DecodeMsgpack func
func (s *PrivateKey) DecodeMsgpack(dec *msgpack.Decoder) error {
	var raw []byte
	err := dec.Decode(&raw)
	if err != nil {
		slog.Warnf("PrivateKey.DecodeMsgpack decode str %v", err)
		return err
	}
	sk, err := x509.ParsePKCS1PrivateKey(raw)
	if err != nil {
		slog.Warnf("PrivateKey.DecodeMsgpack parse pkcs key %v", err)
		return err
	}
	s.PublicKey = sk.PublicKey
	s.D = sk.D
	s.Primes = sk.Primes
	s.Precomputed = sk.Precomputed
	return nil
}

// Decrypt func
func (s *PrivateKey) Decrypt(crypt []byte) (plain []byte, err error) {
	for {
		if len(crypt) <= decryptBlockSize {
			bb, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, (*rsa.PrivateKey)(s), crypt, nil)
			if err != nil {
				return nil, err
			}
			return append(plain, bb...), err
		}
		bb, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, (*rsa.PrivateKey)(s), crypt[:decryptBlockSize], nil)
		if err != nil {
			return nil, err
		}
		plain = append(plain, bb...)
		crypt = crypt[decryptBlockSize:]
	}
}

// MarshalJSON func
func (s *PrivateKey) MarshalJSON() ([]byte, error) {
	return []byte("\"" + base64.RawURLEncoding.EncodeToString(s.Raw()) + "\""), nil
}

// UnmarshalJSON func
func (s *PrivateKey) UnmarshalJSON(b []byte) error {
	if b[0] != '"' && b[len(b)-1] != '"' {
		return fmt.Errorf("Invalid key")
	}
	skey := strings.Replace(string(b[1:len(b)-1]), "\\n", "", -1)
	bb, err := base64.RawURLEncoding.DecodeString(skey)
	if err != nil {
		return fmt.Errorf("Error parse privateKey %v [%s]", err, skey)
	}
	k, err := x509.ParsePKCS1PrivateKey(bb)
	if err != nil {
		return fmt.Errorf("Error parse privateKey %v [%s]", err, string(bb))
	}
	s.D = k.D
	s.E = k.E
	s.N = k.N
	s.Primes = k.Primes
	return ((*rsa.PrivateKey)(s)).Validate()
}

func (s *PrivateKey) String() string {
	bb := x509.MarshalPKCS1PrivateKey((*rsa.PrivateKey)(s))
	return base64.RawURLEncoding.EncodeToString(bb)
}

// Public func
func (s *PrivateKey) Public() *PublicKey {
	rp := (*rsa.PrivateKey)(s).PublicKey
	return (*PublicKey)(&rp)
}

// EncryptKey func
func (s *PrivateKey) EncryptKey(key *PublicKey) EncryptedPrivateKey {
	bb, err := key.Encrypt(s.Raw())
	if err != nil {
		panic(err)
	}
	return (EncryptedPrivateKey)(bb)
}

// Raw func
func (s *PrivateKey) Raw() []byte {
	return x509.MarshalPKCS1PrivateKey((*rsa.PrivateKey)(s))
}

// Equal func
func (s *PrivateKey) Equal(key *PrivateKey) bool {
	if key == nil {
		return s == nil
	}
	k1 := (*rsa.PrivateKey)(s)
	k2 := (*rsa.PrivateKey)(key)
	return k1.E == k2.E && k1.D.Cmp(k2.D) == 0 && k1.N.Cmp(k2.N) == 0
}

var _ msgpack.CustomEncoder = (*PublicKey)(nil)
var _ msgpack.CustomDecoder = (*PublicKey)(nil)

// EncodeMsgpack func
func (s *PublicKey) EncodeMsgpack(enc *msgpack.Encoder) error {
	return enc.Encode(s.Raw())
}

// DecodeMsgpack func
func (s *PublicKey) DecodeMsgpack(dec *msgpack.Decoder) error {
	var raw []byte
	err := dec.Decode(&raw)
	if err != nil {
		return err
	}
	sko, err := x509.ParsePKIXPublicKey(raw)
	if err != nil {
		return err
	}
	if sk, ok := sko.(*rsa.PublicKey); ok {
		s.N = sk.N
		s.E = sk.E
		return nil
	}
	return fmt.Errorf("Unkown type %s", reflect.TypeOf(sko))
}

// Encrypt func
func (s *PublicKey) Encrypt(data []byte) (crypt []byte, err error) {
	for {
		if len(data) <= encryptBlockSize {
			bb, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, (*rsa.PublicKey)(s), data, nil)
			if err != nil {
				return nil, err
			}
			return append(crypt, bb...), err
		}
		bb, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, (*rsa.PublicKey)(s), data[:encryptBlockSize], nil)
		if err != nil {
			return nil, err
		}
		crypt = append(crypt, bb...)
		data = data[encryptBlockSize:]
	}
}

// MarshalJSON func
func (s *PublicKey) MarshalJSON() ([]byte, error) {
	return []byte("\"" + base64.RawURLEncoding.EncodeToString(s.Raw()) + "\""), nil
}

// UnmarshalJSON func
func (s *PublicKey) UnmarshalJSON(b []byte) error {
	if b[0] != '"' && b[len(b)-1] != '"' {
		return fmt.Errorf("Invalid key")
	}
	skey := strings.Replace(string(b[1:len(b)-1]), "\\n", "", -1)
	bb, err := base64.RawURLEncoding.DecodeString(skey)
	if err != nil {
		return fmt.Errorf("Error parse privateKey %v [%s]", err, skey)
	}
	k, err := x509.ParsePKIXPublicKey(bb)
	if err != nil {
		return fmt.Errorf("Error parse privateKey %v [%s]", err, string(bb))
	}
	if pk, ok := k.(*rsa.PublicKey); ok {
		s.E = pk.E
		s.N = pk.N
		return nil
	}
	return fmt.Errorf("Invalid type %s", reflect.TypeOf(k))
}

func (s *PublicKey) String() string {
	return base64.RawURLEncoding.EncodeToString(s.Raw())
}

// Raw func
func (s *PublicKey) Raw() []byte {
	bb, err := x509.MarshalPKIXPublicKey((*rsa.PublicKey)(s))
	if err != nil {
		panic(err)
	}
	return bb
}

// Equal func
func (s *PublicKey) Equal(key *PublicKey) bool {
	if key == nil {
		return s == nil
	}
	k1 := (*rsa.PublicKey)(s)
	k2 := (*rsa.PublicKey)(key)
	return k1.E == k2.E && k1.N.Cmp(k2.N) == 0
}