package asymm

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

	"github.com/emmansun/gmsm/sm2"
	"github.com/emmansun/gmsm/smx509"
)

func GenerateSM2Key() (*sm2.PrivateKey, *ecdsa.PublicKey, error) {
	priv, err := sm2.GenerateKey(rand.Reader)
	if err != nil {
		return nil, nil, err
	}
	return priv, &priv.PublicKey, nil
}

func EncodeSM2PrivateKey(private *sm2.PrivateKey, secret string) ([]byte, error) {
	der, err := smx509.MarshalSM2PrivateKey(private)
	if err != nil {
		return nil, err
	}
	if secret == "" {
		return pem.EncodeToMemory(&pem.Block{Type: "SM2 PRIVATE KEY", Bytes: der}), nil
	}
	blk, err := smx509.EncryptPEMBlock(rand.Reader, "SM2 PRIVATE KEY", der, []byte(secret), smx509.PEMCipherAES256)
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(blk), nil
}

func EncodeSM2PublicKey(public *ecdsa.PublicKey) ([]byte, error) {
	der, err := smx509.MarshalPKIXPublicKey(public)
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: der}), nil
}

func DecodeSM2PrivateKey(private []byte, password string) (*sm2.PrivateKey, error) {
	blk, _ := pem.Decode(private)
	if blk == nil {
		return nil, errors.New("private key error")
	}

	bytes := blk.Bytes
	if smx509.IsEncryptedPEMBlock(blk) {
		if password == "" {
			return nil, errors.New("private key is encrypted but password is empty")
		}
		var err error
		bytes, err = smx509.DecryptPEMBlock(blk, []byte(password))
		if err != nil {
			return nil, err
		}
	}

	if key, err := smx509.ParseSM2PrivateKey(bytes); err == nil {
		return key, nil
	}
	pkcs8, err := smx509.ParsePKCS8PrivateKey(bytes)
	if err != nil {
		return nil, err
	}
	key, ok := pkcs8.(*sm2.PrivateKey)
	if !ok {
		return nil, errors.New("private key is not SM2")
	}
	return key, nil
}

func DecodeSM2PublicKey(pubStr []byte) (*ecdsa.PublicKey, error) {
	blk, _ := pem.Decode(pubStr)
	if blk == nil {
		return nil, errors.New("public key error")
	}
	pubAny, err := smx509.ParsePKIXPublicKey(blk.Bytes)
	if err != nil {
		return nil, err
	}
	pub, ok := pubAny.(*ecdsa.PublicKey)
	if !ok {
		return nil, errors.New("public key is not ECDSA/SM2")
	}
	if !sm2.IsSM2PublicKey(pub) {
		return nil, errors.New("public key is not SM2")
	}
	return pub, nil
}

func SM2EncryptASN1(pub *ecdsa.PublicKey, data []byte) ([]byte, error) {
	return sm2.EncryptASN1(rand.Reader, pub, data)
}

func SM2DecryptASN1(priv *sm2.PrivateKey, data []byte) ([]byte, error) {
	return priv.Decrypt(nil, data, sm2.ASN1DecrypterOpts)
}

func SM2Sign(priv *sm2.PrivateKey, msg, uid []byte) ([]byte, error) {
	if len(uid) == 0 {
		uid = nil
	}
	return sm2.SignASN1(rand.Reader, priv, msg, sm2.NewSM2SignerOption(true, uid))
}

func SM2Verify(pub *ecdsa.PublicKey, msg, sig, uid []byte) bool {
	if len(uid) == 0 {
		uid = nil
	}
	return sm2.VerifyASN1WithSM2(pub, uid, msg, sig)
}

func SM2SignByPEM(msg, priKey []byte, password string, uid []byte) ([]byte, error) {
	priv, err := DecodeSM2PrivateKey(priKey, password)
	if err != nil {
		return nil, err
	}
	return SM2Sign(priv, msg, uid)
}

func SM2VerifyByPEM(sig, msg, pubKey []byte, uid []byte) (bool, error) {
	pub, err := DecodeSM2PublicKey(pubKey)
	if err != nil {
		return false, err
	}
	return SM2Verify(pub, msg, sig, uid), nil
}

func IsSM2PublicKey(public crypto.PublicKey) bool {
	return sm2.IsSM2PublicKey(public)
}