package asymm

import (
	"crypto/rand"
	"encoding/pem"
	"errors"

	gmsm3 "github.com/emmansun/gmsm/sm3"
	gmsm9 "github.com/emmansun/gmsm/sm9"
)

const (
	SM9SignHID    byte = 0x01
	SM9EncryptHID byte = 0x03
)

func GenerateSM9SignMasterKey() (*gmsm9.SignMasterPrivateKey, *gmsm9.SignMasterPublicKey, error) {
	priv, err := gmsm9.GenerateSignMasterKey(rand.Reader)
	if err != nil {
		return nil, nil, err
	}
	return priv, priv.PublicKey(), nil
}

func GenerateSM9EncryptMasterKey() (*gmsm9.EncryptMasterPrivateKey, *gmsm9.EncryptMasterPublicKey, error) {
	priv, err := gmsm9.GenerateEncryptMasterKey(rand.Reader)
	if err != nil {
		return nil, nil, err
	}
	return priv, priv.PublicKey(), nil
}

func GenerateSM9SignUserKey(master *gmsm9.SignMasterPrivateKey, uid []byte, hid byte) (*gmsm9.SignPrivateKey, error) {
	if master == nil {
		return nil, errors.New("sm9 sign master key is nil")
	}
	if hid == 0 {
		hid = SM9SignHID
	}
	return master.GenerateUserKey(uid, hid)
}

func GenerateSM9EncryptUserKey(master *gmsm9.EncryptMasterPrivateKey, uid []byte, hid byte) (*gmsm9.EncryptPrivateKey, error) {
	if master == nil {
		return nil, errors.New("sm9 encrypt master key is nil")
	}
	if hid == 0 {
		hid = SM9EncryptHID
	}
	return master.GenerateUserKey(uid, hid)
}

func EncodeSM9SignMasterPrivateKey(key *gmsm9.SignMasterPrivateKey) ([]byte, error) {
	if key == nil {
		return nil, errors.New("sm9 sign master private key is nil")
	}
	der, err := key.MarshalASN1()
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(&pem.Block{Type: "SM9 SIGN MASTER PRIVATE KEY", Bytes: der}), nil
}

func DecodeSM9SignMasterPrivateKey(data []byte) (*gmsm9.SignMasterPrivateKey, error) {
	der, err := pemOrDER(data)
	if err != nil {
		return nil, err
	}
	return gmsm9.UnmarshalSignMasterPrivateKeyASN1(der)
}

func EncodeSM9SignMasterPublicKey(key *gmsm9.SignMasterPublicKey) ([]byte, error) {
	if key == nil {
		return nil, errors.New("sm9 sign master public key is nil")
	}
	der, err := key.MarshalASN1()
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(&pem.Block{Type: "SM9 SIGN MASTER PUBLIC KEY", Bytes: der}), nil
}

func DecodeSM9SignMasterPublicKey(data []byte) (*gmsm9.SignMasterPublicKey, error) {
	der, err := pemOrDER(data)
	if err != nil {
		return nil, err
	}
	return gmsm9.UnmarshalSignMasterPublicKeyASN1(der)
}

func EncodeSM9SignPrivateKey(key *gmsm9.SignPrivateKey) ([]byte, error) {
	if key == nil {
		return nil, errors.New("sm9 sign private key is nil")
	}
	der, err := key.MarshalASN1()
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(&pem.Block{Type: "SM9 SIGN PRIVATE KEY", Bytes: der}), nil
}

func DecodeSM9SignPrivateKey(data []byte) (*gmsm9.SignPrivateKey, error) {
	der, err := pemOrDER(data)
	if err != nil {
		return nil, err
	}
	return gmsm9.UnmarshalSignPrivateKeyASN1(der)
}

func EncodeSM9EncryptMasterPrivateKey(key *gmsm9.EncryptMasterPrivateKey) ([]byte, error) {
	if key == nil {
		return nil, errors.New("sm9 encrypt master private key is nil")
	}
	der, err := key.MarshalASN1()
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(&pem.Block{Type: "SM9 ENCRYPT MASTER PRIVATE KEY", Bytes: der}), nil
}

func DecodeSM9EncryptMasterPrivateKey(data []byte) (*gmsm9.EncryptMasterPrivateKey, error) {
	der, err := pemOrDER(data)
	if err != nil {
		return nil, err
	}
	return gmsm9.UnmarshalEncryptMasterPrivateKeyASN1(der)
}

func EncodeSM9EncryptMasterPublicKey(key *gmsm9.EncryptMasterPublicKey) ([]byte, error) {
	if key == nil {
		return nil, errors.New("sm9 encrypt master public key is nil")
	}
	der, err := key.MarshalASN1()
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(&pem.Block{Type: "SM9 ENCRYPT MASTER PUBLIC KEY", Bytes: der}), nil
}

func DecodeSM9EncryptMasterPublicKey(data []byte) (*gmsm9.EncryptMasterPublicKey, error) {
	der, err := pemOrDER(data)
	if err != nil {
		return nil, err
	}
	return gmsm9.UnmarshalEncryptMasterPublicKeyASN1(der)
}

func EncodeSM9EncryptPrivateKey(key *gmsm9.EncryptPrivateKey) ([]byte, error) {
	if key == nil {
		return nil, errors.New("sm9 encrypt private key is nil")
	}
	der, err := key.MarshalASN1()
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(&pem.Block{Type: "SM9 ENCRYPT PRIVATE KEY", Bytes: der}), nil
}

func DecodeSM9EncryptPrivateKey(data []byte) (*gmsm9.EncryptPrivateKey, error) {
	der, err := pemOrDER(data)
	if err != nil {
		return nil, err
	}
	return gmsm9.UnmarshalEncryptPrivateKeyASN1(der)
}

func SM9SignHashASN1(priv *gmsm9.SignPrivateKey, hash []byte) ([]byte, error) {
	if priv == nil {
		return nil, errors.New("sm9 sign private key is nil")
	}
	return gmsm9.SignASN1(rand.Reader, priv, hash)
}

func SM9SignASN1(priv *gmsm9.SignPrivateKey, message []byte) ([]byte, error) {
	sum := gmsm3.Sum(message)
	return SM9SignHashASN1(priv, sum[:])
}

func SM9VerifyHashASN1(pub *gmsm9.SignMasterPublicKey, uid []byte, hid byte, hash, sig []byte) bool {
	if pub == nil {
		return false
	}
	if hid == 0 {
		hid = SM9SignHID
	}
	return gmsm9.VerifyASN1(pub, uid, hid, hash, sig)
}

func SM9VerifyASN1(pub *gmsm9.SignMasterPublicKey, uid []byte, hid byte, message, sig []byte) bool {
	sum := gmsm3.Sum(message)
	return SM9VerifyHashASN1(pub, uid, hid, sum[:], sig)
}

func SM9Encrypt(pub *gmsm9.EncryptMasterPublicKey, uid []byte, hid byte, plaintext []byte) ([]byte, error) {
	if pub == nil {
		return nil, errors.New("sm9 encrypt master public key is nil")
	}
	if hid == 0 {
		hid = SM9EncryptHID
	}
	return gmsm9.Encrypt(rand.Reader, pub, uid, hid, plaintext, gmsm9.SM4CBCEncrypterOpts)
}

func SM9Decrypt(priv *gmsm9.EncryptPrivateKey, uid, ciphertext []byte) ([]byte, error) {
	if priv == nil {
		return nil, errors.New("sm9 encrypt private key is nil")
	}
	return gmsm9.Decrypt(priv, uid, ciphertext, gmsm9.SM4CBCEncrypterOpts)
}

func SM9EncryptASN1(pub *gmsm9.EncryptMasterPublicKey, uid []byte, hid byte, plaintext []byte) ([]byte, error) {
	if pub == nil {
		return nil, errors.New("sm9 encrypt master public key is nil")
	}
	if hid == 0 {
		hid = SM9EncryptHID
	}
	return gmsm9.EncryptASN1(rand.Reader, pub, uid, hid, plaintext, gmsm9.SM4CBCEncrypterOpts)
}

func SM9DecryptASN1(priv *gmsm9.EncryptPrivateKey, uid, ciphertext []byte) ([]byte, error) {
	if priv == nil {
		return nil, errors.New("sm9 encrypt private key is nil")
	}
	return gmsm9.DecryptASN1(priv, uid, ciphertext)
}

func pemOrDER(data []byte) ([]byte, error) {
	if len(data) == 0 {
		return nil, errors.New("empty key data")
	}
	if blk, _ := pem.Decode(data); blk != nil {
		return blk.Bytes, nil
	}
	return data, nil
}