package pbkd

import (
	"encoding/asn1"

	"golang.org/x/crypto/cryptobyte"
	cryptobyte_asn1 "golang.org/x/crypto/cryptobyte/asn1"
	"xdx.jelly/xgcl/gerrors"
)

//	 id-PBES2 OBJECT IDENTIFIER ::= {pkcs-5 13}
//
// The parameters field associated with this OID in an
// AlgorithmIdentifier shall have type PBES2-params:
//
//	PBES2-params ::= SEQUENCE {
//	 keyDerivationFunc AlgorithmIdentifier {{PBES2-KDFs}},
//	 encryptionScheme AlgorithmIdentifier {{PBES2-Encs}}
//	}
//
//	PBKDF2-params ::= SEQUENCE {
//		salt CHOICE {
//			specified OCTET STRING,
//			otherSource AlgorithmIdentifier {{PBKDF2-SaltSources}}
//		},
//		iterationCount INTEGER (1..MAX),
//		keyLength INTEGER (1..MAX) OPTIONAL,
//		prf AlgorithmIdentifier {{PBKDF2-PRFs}} DEFAULT algid-hmacWithSHA1 }

var (
	// oidHmacWithSHA1 = asn1.ObjectIdentifier{1, 2, 156, 10197, 1, 401, 2}
	IdHmacWithSm3 = asn1.ObjectIdentifier{1, 2, 156, 10197, 1, 401, 2} //SM3密码杂凑算法，有密钥使用
	IdEncSM4CBC   = asn1.ObjectIdentifier{1, 2, 156, 10197, 1, 104, 2} // SM4-CBC

	IdPkcs5PBKDF2 = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 5, 12} // pkcs5PBKDF2 (PKCS #5 v2.0)
	IdPkcs5PBES2  = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 5, 13} // pkcs5PBES2 (PKCS #5 v2.0)
)

func MarshalASN1PBKDF2(salt []byte, iterationCount int, keyLength int) ([]byte, error) {
	pbkdf2 := &PBKDF2{
		ObjectIdentifier: IdPkcs5PBKDF2,
	}
	pbkdf2.Param.Salt = salt
	pbkdf2.Param.IterationCount = iterationCount
	pbkdf2.Param.KeyLength = new(int)
	*pbkdf2.Param.KeyLength = keyLength
	pbkdf2.Param.Prf.ObjectIdentifier = IdHmacWithSm3
	return pbkdf2.MarshalASN1()
}

type PBKDF2 struct {
	// Salt only support OCTET STRING
	asn1.ObjectIdentifier
	Param struct {
		Salt           []byte
		IterationCount int
		KeyLength      *int
		Prf            struct {
			asn1.ObjectIdentifier // always be idHmacWithSm3
		}
	}
}

func (p *PBKDF2) MarshalASN1() ([]byte, error) {
	var b cryptobyte.Builder
	b.AddASN1(cryptobyte_asn1.SEQUENCE, func(b *cryptobyte.Builder) {
		b.AddASN1OctetString(p.Param.Salt)
		b.AddASN1Int64(int64(p.Param.IterationCount))
		if p.Param.KeyLength != nil {
			b.AddASN1Int64(int64(*p.Param.KeyLength))
		}
		b.AddASN1(cryptobyte_asn1.SEQUENCE, func(b *cryptobyte.Builder) {
			b.AddASN1ObjectIdentifier(p.Param.Prf.ObjectIdentifier)
		})
	})
	return b.Bytes()
}

func (p *PBKDF2) UnmarshalASN1(data []byte) (rest []byte, err error) {
	input := cryptobyte.String(data)
	var inner cryptobyte.String
	var paramString cryptobyte.String
	var algIdString cryptobyte.String

	if !input.ReadASN1(&inner, cryptobyte_asn1.SEQUENCE) ||
		!inner.ReadASN1ObjectIdentifier(&p.ObjectIdentifier) ||
		!inner.ReadASN1(&paramString, cryptobyte_asn1.SEQUENCE) ||
		!inner.Empty() ||
		!paramString.ReadASN1Bytes(&p.Param.Salt, cryptobyte_asn1.OCTET_STRING) ||
		!paramString.ReadASN1Integer(&p.Param.IterationCount) {
		return data, gerrors.WithAnnotating(ErrInvalidInput, "read ASN.1 SEQUENCE failed")
	}
	if paramString.PeekASN1Tag(cryptobyte_asn1.INTEGER) {
		keyLength := 0
		if !paramString.ReadASN1Integer(&keyLength) {
			return data, gerrors.WithAnnotating(ErrInvalidInput, "read ASN.1 SEQUENCE failed")
		}
		if p.Param.KeyLength == nil {
			p.Param.KeyLength = new(int)
		}
		*p.Param.KeyLength = keyLength
	}
	if !paramString.ReadASN1(&algIdString, cryptobyte_asn1.SEQUENCE) ||
		!paramString.Empty() ||
		!algIdString.ReadASN1ObjectIdentifier(&p.Param.Prf.ObjectIdentifier) ||
		!algIdString.Empty() {
		return data, gerrors.WithAnnotating(ErrInvalidInput, "read ASN.1 SEQUENCE failed")
	}
	return input, nil
}

type PBES2 struct {
	asn1.ObjectIdentifier
	Param struct {
		KeyDerivationFunc PBKDF2
		EncryptionScheme  struct {
			asn1.ObjectIdentifier        // 1.2.156.10197.1.104.2 (SM4_CBC, GmSSL, but not the same with GM/T 0091)
			iv                    []byte // The IV, always 16 bytes
		}
	}
}

func (p *PBES2) Salt() []byte {
	return p.Param.KeyDerivationFunc.Param.Salt
}

func (p *PBES2) IterationCount() int {
	return p.Param.KeyDerivationFunc.Param.IterationCount
}

func (p *PBES2) IV() []byte {
	return p.Param.EncryptionScheme.iv
}

func (p *PBES2) MarshalASN1() ([]byte, error) {
	var b cryptobyte.Builder
	b.AddASN1(cryptobyte_asn1.SEQUENCE, func(b *cryptobyte.Builder) {
		b.AddASN1ObjectIdentifier(p.ObjectIdentifier)
		b.AddASN1(cryptobyte_asn1.SEQUENCE, func(b *cryptobyte.Builder) {
			b.MarshalASN1(&p.Param.KeyDerivationFunc)
			b.AddASN1(cryptobyte_asn1.SEQUENCE, func(b *cryptobyte.Builder) {
				b.AddASN1ObjectIdentifier(p.Param.EncryptionScheme.ObjectIdentifier)
				b.AddASN1OctetString(p.Param.EncryptionScheme.iv)
			})
		})
	})
	return b.Bytes()
}

func (p *PBES2) UnmarshalASN1(data []byte) ([]byte, error) {
	input := cryptobyte.String(data)
	var inner cryptobyte.String
	var paramInner cryptobyte.String
	var encryptionSchemeString cryptobyte.String

	if !input.ReadASN1(&inner, cryptobyte_asn1.SEQUENCE) ||
		!inner.ReadASN1ObjectIdentifier(&p.ObjectIdentifier) ||
		!inner.ReadASN1(&paramInner, cryptobyte_asn1.SEQUENCE) ||
		!inner.Empty() {
		return data, gerrors.WithAnnotating(ErrInvalidInput, "read ASN.1 SEQUENCE failed")
	}

	rest, err := p.Param.KeyDerivationFunc.UnmarshalASN1(paramInner)
	if err != nil {
		return data, gerrors.WithAnnotating(ErrInvalidInput, "read ASN.1 SEQUENCE failed")
	}
	encryptionSchemeString = cryptobyte.String(rest)
	if !encryptionSchemeString.ReadASN1(&inner, cryptobyte_asn1.SEQUENCE) ||
		!encryptionSchemeString.Empty() ||
		!inner.ReadASN1ObjectIdentifier(&p.Param.EncryptionScheme.ObjectIdentifier) ||
		!inner.ReadASN1Bytes(&p.Param.EncryptionScheme.iv, cryptobyte_asn1.OCTET_STRING) ||
		!inner.Empty() {
		return data, gerrors.WithAnnotating(ErrInvalidInput, "read ASN.1 SEQUENCE failed")
	}
	return input, nil
}