package util

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/elliptic"
	"crypto/hmac"
	"crypto/md5"
	"crypto/rand"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"crypto/x509/pkix"
	"encoding/asn1"
	"fmt"
	"hash"
	"math/big"
	"reflect"
	"voucher/internal/pkg/cmb/sm2/model"
)

/*
 * reference to RFC5959 and RFC2898
 */

var (
	oidPBES2  = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 5, 13} // id-PBES2(PBES2)
	oidPBKDF2 = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 5, 12} // id-PBKDF2

	oidAES128CBC = asn1.ObjectIdentifier{2, 16, 840, 1, 101, 3, 4, 1, 2}
	oidAES256CBC = asn1.ObjectIdentifier{2, 16, 840, 1, 101, 3, 4, 1, 42}

	oidKEYMD5    = asn1.ObjectIdentifier{1, 2, 840, 113549, 2, 5}
	oidKEYSHA1   = asn1.ObjectIdentifier{1, 2, 840, 113549, 2, 7}
	oidKEYSHA256 = asn1.ObjectIdentifier{1, 2, 840, 113549, 2, 9}
	oidKEYSHA512 = asn1.ObjectIdentifier{1, 2, 840, 113549, 2, 11}

	oidSM2 = asn1.ObjectIdentifier{1, 2, 840, 10045, 2, 1}
)

type Sm2PrivateKey struct {
	Version       int
	PrivateKey    []byte
	NamedCurveOID asn1.ObjectIdentifier `asn1:"optional,explicit,tag:0"`
	PublicKey     asn1.BitString        `asn1:"optional,explicit,tag:1"`
}

type pkcs8 struct {
	Version    int
	Algo       pkix.AlgorithmIdentifier
	PrivateKey []byte
}

// EncryptedPrivateKeyInfo reference to https://www.rfc-editor.org/rfc/rfc5958.txt
type EncryptedPrivateKeyInfo struct {
	EncryptionAlgorithm Pbes2Algorithms
	EncryptedData       []byte
}

// Pbes2Algorithms reference to https://www.ietf.org/rfc/rfc2898.txt
type Pbes2Algorithms struct {
	IdPBES2     asn1.ObjectIdentifier
	Pbes2Params Pbes2Params
}

// Pbes2Params reference to https://www.ietf.org/rfc/rfc2898.txt
type Pbes2Params struct {
	KeyDerivationFunc Pbes2KDfs // PBES2-KDFs
	EncryptionScheme  Pbes2Encs // PBES2-Encs
}

// Pbes2KDfs reference to https://www.ietf.org/rfc/rfc2898.txt
type Pbes2KDfs struct {
	IdPBKDF2    asn1.ObjectIdentifier
	Pkdf2Params Pkdf2Params
}
type Pbes2Encs struct {
	EncryAlgo asn1.ObjectIdentifier
	IV        []byte
}

// Pkdf2Params reference to https://www.ietf.org/rfc/rfc2898.txt
type Pkdf2Params struct {
	Salt           []byte
	IterationCount int
	Prf            pkix.AlgorithmIdentifier
}

func ParsePrivateKey(bytes []byte, pwd []byte) (*model.PrivateKey, error) {
	var priKey Sm2PrivateKey
	var err error
	if pwd == nil {
		priKey, err = ParsePKCS8UnEncryptedPrivateKey(bytes)
	} else {
		priKey, err = ParsePKCS8EncryptedPrivateKey(bytes, pwd)
	}
	if err != nil {
		return nil, fmt.Errorf("parse private key err: %s", err.Error())
	}
	curve := NewP256Sm2()
	k := new(big.Int).SetBytes(priKey.PrivateKey)
	curveOrder := curve.Params().N
	if k.Cmp(curveOrder) >= 0 {
		return nil, fmt.Errorf("invalid elliptic curve private key value")
	}
	privateKey := make([]byte, (curveOrder.BitLen()+7)/8)
	for len(priKey.PrivateKey) > len(privateKey) {
		if priKey.PrivateKey[0] != 0 {
			return nil, fmt.Errorf("invalid private key")
		}
		priKey.PrivateKey = priKey.PrivateKey[1:]
	}
	copy(privateKey[len(privateKey)-len(priKey.PrivateKey):], priKey.PrivateKey)
	x, y := curve.ScalarBaseMult(privateKey)
	return &model.PrivateKey{
		PublicKey: &model.PublicKey{
			Curve: curve,
			X:     x,
			Y:     y,
		},
		D: k,
	}, nil
}

func ParsePKCS8UnEncryptedPrivateKey(bytes []byte) (Sm2PrivateKey, error) {
	var pk pkcs8
	var priKey Sm2PrivateKey

	if _, err := asn1.Unmarshal(bytes, &pk); err != nil {
		return priKey, err
	}
	if !reflect.DeepEqual(pk.Algo.Algorithm, oidSM2) {
		return priKey, fmt.Errorf("not sm2 elliptic curve")
	}

	if _, err := asn1.Unmarshal(pk.PrivateKey, &priKey); err != nil {
		return priKey, fmt.Errorf("privateKey is not sm2 private key: %s", err.Error())
	}
	return priKey, nil
}

func ParsePKCS8EncryptedPrivateKey(bytes, pwd []byte) (Sm2PrivateKey, error) {
	var keyInfo EncryptedPrivateKeyInfo
	var priKey Sm2PrivateKey
	_, err := asn1.Unmarshal(bytes, &keyInfo)
	if err != nil {
		return priKey, fmt.Errorf("privateKey is not sm2 encrypted private key: %s", err.Error())
	}
	if !reflect.DeepEqual(keyInfo.EncryptionAlgorithm.IdPBES2, oidPBES2) {
		return priKey, fmt.Errorf("x509: only support PBES2")
	}
	encryptionScheme := keyInfo.EncryptionAlgorithm.Pbes2Params.EncryptionScheme
	keyDerivationFunc := keyInfo.EncryptionAlgorithm.Pbes2Params.KeyDerivationFunc
	if !reflect.DeepEqual(keyDerivationFunc.IdPBKDF2, oidPBKDF2) {
		return priKey, fmt.Errorf("x509: only support PBKDF2")
	}
	pkdf2Params := keyDerivationFunc.Pkdf2Params
	if !reflect.DeepEqual(encryptionScheme.EncryAlgo, oidAES128CBC) &&
		!reflect.DeepEqual(encryptionScheme.EncryAlgo, oidAES256CBC) {
		return priKey, fmt.Errorf("x509: only support AES")
	}
	iv := encryptionScheme.IV
	salt := pkdf2Params.Salt
	iter := pkdf2Params.IterationCount
	encryptedKey := keyInfo.EncryptedData
	var key []byte
	switch {
	case pkdf2Params.Prf.Algorithm.Equal(oidKEYMD5):
		key = pbkdf(pwd, salt, iter, 32, md5.New)
		break
	case pkdf2Params.Prf.Algorithm.Equal(oidKEYSHA1):
		key = pbkdf(pwd, salt, iter, 32, sha1.New)
		break
	case pkdf2Params.Prf.Algorithm.Equal(oidKEYSHA256):
		key = pbkdf(pwd, salt, iter, 32, sha256.New)
		break
	case pkdf2Params.Prf.Algorithm.Equal(oidKEYSHA512):
		key = pbkdf(pwd, salt, iter, 32, sha512.New)
		break
	default:
		return priKey, fmt.Errorf("x509: unknown hash algorithm")
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return priKey, err
	}
	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(encryptedKey, encryptedKey)
	return ParsePKCS8UnEncryptedPrivateKey(encryptedKey)
}

// copy from crypto/pbkdf2.go
func pbkdf(password, salt []byte, iter, keyLen int, h func() hash.Hash) []byte {
	prf := hmac.New(h, password)
	hashLen := prf.Size()
	numBlocks := (keyLen + hashLen - 1) / hashLen

	var buf [4]byte
	dk := make([]byte, 0, numBlocks*hashLen)
	U := make([]byte, hashLen)
	for block := 1; block <= numBlocks; block++ {
		// N.B.: || means concatenation, ^ means XOR
		// for each block T_i = U_1 ^ U_2 ^ ... ^ U_iter
		// U_1 = PRF(password, salt || uint(i))
		prf.Reset()
		prf.Write(salt)
		buf[0] = byte(block >> 24)
		buf[1] = byte(block >> 16)
		buf[2] = byte(block >> 8)
		buf[3] = byte(block)
		prf.Write(buf[:4])
		dk = prf.Sum(dk)
		T := dk[len(dk)-hashLen:]
		copy(U, T)

		// U_n = PRF(password, U_(n-1))
		for n := 2; n <= iter; n++ {
			prf.Reset()
			prf.Write(U)
			U = U[:0]
			U = prf.Sum(U)
			for x := range U {
				T[x] ^= U[x]
			}
		}
	}
	return dk[:keyLen]
}

func MarshalSm2PrivateKey(key *model.PrivateKey, pwd []byte) ([]byte, error) {
	if pwd == nil {
		return MarshalSm2UnEncryptedPrivateKey(key)
	}
	return MarshalSm2EncryptedPrivateKey(key, pwd)
}

func MarshalSm2EncryptedPrivateKey(priKey *model.PrivateKey, pwd []byte) ([]byte, error) {
	der, err := MarshalSm2UnEncryptedPrivateKey(priKey)
	if err != nil {
		return nil, err
	}
	iter := 2048
	salt := make([]byte, 8)
	iv := make([]byte, 16)
	rand.Reader.Read(salt)
	rand.Reader.Read(iv)
	key := pbkdf(pwd, salt, iter, 32, sha1.New) // 默认是SHA1
	padding := aes.BlockSize - len(der)%aes.BlockSize
	if padding > 0 {
		n := len(der)
		der = append(der, make([]byte, padding)...)
		for i := 0; i < padding; i++ {
			der[n+i] = byte(padding)
		}
	}
	encryptedKey := make([]byte, len(der))
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(encryptedKey, der)
	var algorithmIdentifier pkix.AlgorithmIdentifier
	algorithmIdentifier.Algorithm = oidKEYSHA1
	algorithmIdentifier.Parameters.Tag = 5
	algorithmIdentifier.Parameters.IsCompound = false
	algorithmIdentifier.Parameters.FullBytes = []byte{5, 0}
	keyDerivationFunc := Pbes2KDfs{
		oidPBKDF2,
		Pkdf2Params{
			salt,
			iter,
			algorithmIdentifier,
		},
	}
	encryptionScheme := Pbes2Encs{
		oidAES256CBC,
		iv,
	}
	pbes2Algorithms := Pbes2Algorithms{
		oidPBES2,
		Pbes2Params{
			keyDerivationFunc,
			encryptionScheme,
		},
	}
	encryptedPkey := EncryptedPrivateKeyInfo{
		pbes2Algorithms,
		encryptedKey,
	}
	return asn1.Marshal(encryptedPkey)
}

func MarshalSm2UnEncryptedPrivateKey(key *model.PrivateKey) ([]byte, error) {
	var r pkcs8
	var pri Sm2PrivateKey
	var algo pkix.AlgorithmIdentifier

	algo.Algorithm = oidSM2
	algo.Parameters.Class = 0
	algo.Parameters.Tag = 6
	algo.Parameters.IsCompound = false
	algo.Parameters.FullBytes = []byte{6, 8, 42, 129, 28, 207, 85, 1, 130, 45} // asn1.Marshal(asn1.ObjectIdentifier{1, 2, 156, 10197, 1, 301})
	pri.Version = 1
	pri.NamedCurveOID = oidNamedCurveP256SM2
	pri.PublicKey = asn1.BitString{Bytes: elliptic.Marshal(key.Curve, key.X, key.Y)}
	pri.PrivateKey = key.D.Bytes()
	r.Version = 0
	r.Algo = algo
	r.PrivateKey, _ = asn1.Marshal(pri)
	return asn1.Marshal(r)
}

func MarshalSm2PublicKey(key *model.PublicKey) ([]byte, error) {
	var r PKIXPublicKey
	var algo pkix.AlgorithmIdentifier

	if key.Curve.Params() != NewP256Sm2().Params() {
		return nil, fmt.Errorf("x509: unsupported elliptic curve")
	}
	algo.Algorithm = oidSM2
	algo.Parameters.Class = 0
	algo.Parameters.Tag = 6
	algo.Parameters.IsCompound = false
	algo.Parameters.FullBytes = []byte{6, 8, 42, 129, 28, 207, 85, 1, 130, 45} // asn1.Marshal(asn1.ObjectIdentifier{1, 2, 156, 10197, 1, 301})
	r.Algo = algo
	r.BitString = asn1.BitString{Bytes: elliptic.Marshal(key.Curve, key.X, key.Y)}
	return asn1.Marshal(r)
}