package sm9

import (
	encoding_asn1 "encoding/asn1"
	"math/big"
	"time"

	"golang.org/x/crypto/cryptobyte"
	cryptobyte_asn1 "golang.org/x/crypto/cryptobyte/asn1"
	"xdx.jelly/xgcl/gerrors"
	"xdx.jelly/xgcl/gmath"
	oid "xdx.jelly/xgcl/identifier"
	"xdx.jelly/xgcl/sm/sm9/errors"
	"xdx.jelly/xgcl/sm/sm9/internal/bn256"
)

// GM/T 0080 SM9密码算法使用规范
// ASN.1编解码

// GMT/T 0006 OID定义
var (
	OidSm9               = oid.OIDSM9
	OidSm9Sign           = oid.OIDSM9Signature
	OidSm9KeyExchange    = oid.OIDSM9KeyExchange
	OidSm9Encryption     = oid.OIDSM9Encryption
	OidSm9EncryptionEcb  = oid.OIDSM9EncryptionECB
	OidSm9EncryptionCbc  = oid.OIDSM9EncryptionCBC
	OidSm9EncryptionCfb  = oid.OIDSM9EncryptionCFB
	OidSm9EncryptionOfb  = oid.OIDSM9EncryptionOFB
	OidSM9KeyEncupsulate = oid.OIDSM9KeyEncupsulate
)

const numBytes = 32

// Identifier GM/T 0090定义身份标识
//
// deprecated: Identifier moves to package xdx.jelly/xibc
type Identifier struct {
	// 标准中多了一个explicit
	Version      int `asn1:"default:0"`
	IdentityType encoding_asn1.ObjectIdentifier
	Alias        string `asn1:"utf8"`
	IdentityData []byte
	Serial       *big.Int  `asn1:"optional,explicit,tag:0"`
	ValidStart   time.Time `asn1:"generalized"`
	ValidEnd     time.Time `asn1:"optional,explicit,generalized,tag:1"`
	IdExtentions []byte    `asn1:"optional,explicit,tag:2"`
}

func (id *Identifier) MarshalASN1() ([]byte, error) {
	return encoding_asn1.Marshal(*id)
}

func (id *Identifier) UnmarshalASN1(b []byte) (rest []byte, err error) {
	return encoding_asn1.Unmarshal(b, id)
}

////////////////////////////////////////////////////////////////
//
//						主签名密钥
//
////////////////////////////////////////////////////////////////

func (k *MastSignPrivateKey) MarshalASN1() ([]byte, error) {
	var b cryptobyte.Builder
	b.AddASN1BigInt(&k.Int)
	return b.Bytes()
}

func (k *MastSignPrivateKey) UnmarshalASN1(data []byte) (rest []byte, err error) {
	input := cryptobyte.String(data)
	ok := input.ReadASN1Integer(&k.Int)
	if !ok {
		return data, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse MastSignPrivateKey failed, read ASN.1 Integer error")
	}
	return []byte(input), nil
}

// MarshalASN1 返回ASN.1编码，如果compressed为true，则返回的是压缩形式
//
// 注：不支持混合形式输出
func (k *MastSignPublicKey) MarshalASN1(compressed bool) ([]byte, error) {
	return k.G2.MarshalASN1(compressed)
}

// UnmarshalASN1 ASN.1解码MastSignPublicKey
//
// 注：支持tag为0,2,3,4,6,7即无穷远点(0)，非压缩形式(4)，压缩形式(2,3)，混合形式(6,7)
func (k *MastSignPublicKey) UnmarshalASN1(data []byte) (rest []byte, err error) {
	return k.G2.UnmarshalASN1(data)
}

////////////////////////////////////////////////////////////////
//
//						主加密密钥
//
////////////////////////////////////////////////////////////////

func (k *MastEncPrivateKey) MarshalASN1() ([]byte, error) {
	var b cryptobyte.Builder
	b.AddASN1BigInt(&k.Int)
	return b.Bytes()
}

func (k *MastEncPrivateKey) UnmarshalASN1(data []byte) (rest []byte, err error) {
	input := cryptobyte.String(data)
	ok := input.ReadASN1Integer(&k.Int)
	if !ok {
		return data, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse MastEncPrivateKey failed, read ASN.1 Integer error")
	}
	return []byte(input), nil
}

// MarshalASN1 返回ASN.1编码，如果compressed为true，则返回的是压缩形式
//
// 注：不支持混合形式
func (k *MastEncPublicKey) MarshalASN1(compressed bool) ([]byte, error) {
	return k.G1.MarshalASN1(compressed)
}

// UnmarshalASN1 ASN.1解码MastSignPublicKey
//
// 注：支持tag为0,2,3,4,6,7即无穷远点(0)，非压缩形式(4)，压缩形式(2,3)，混合形式(6,7)
func (k *MastEncPublicKey) UnmarshalASN1(data []byte) (rest []byte, err error) {
	return k.G1.UnmarshalASN1(data)
}

////////////////////////////////////////////////////////////////
//
//						用户密钥
//
////////////////////////////////////////////////////////////////

func (u *UserSignKey) MarshalASN1(compressed bool) ([]byte, error) {
	return u.G1.MarshalASN1(compressed)
}

func (u *UserSignKey) UnmarshalASN1(data []byte) (rest []byte, err error) {
	return u.G1.UnmarshalASN1(data)
}

func (u *UserEncKey) MarshalASN1(compressed bool) ([]byte, error) {
	return u.G2.MarshalASN1(compressed)

}

func (u *UserEncKey) UnmarshalASN1(data []byte) (rest []byte, err error) {
	return u.G2.UnmarshalASN1(data)
}

////////////////////////////////////////////////////////////////
//
//						签名结构
//
////////////////////////////////////////////////////////////////

// SM9Signature ::= SEQUENCE{
// 	H	OCTET STRING,	——杂凑分量，算法是H2（见GB/T 38635.2—2020）
// 	S	BIT STRING   	——签名结果（见GB/T 38635.2—2020）
// }

func (s *Signature) MarshalASN1(compressed bool) ([]byte, error) {
	var b cryptobyte.Builder
	h := make([]byte, numBytes)
	err := gmath.FillBytes(&s.H, h)
	if err != nil {
		return nil, gerrors.WithAnnotating(
			gerrors.ChainErrors(errors.ErrDecodeASN1Failed, err),
			"Signature.H is bigger than 256 bits",
		)
	}

	b.AddASN1(cryptobyte_asn1.SEQUENCE, func(b *cryptobyte.Builder) {
		b.AddASN1OctetString(h)
		b.AddASN1BitString(bn256.MarshalG1(&s.S, compressed))
	})

	return b.Bytes()
}

func (s *Signature) UnmarshalASN1(data []byte) (rest []byte, err error) {
	input := cryptobyte.String(data)
	var inner, hBytes cryptobyte.String
	var sBitString encoding_asn1.BitString
	if !input.ReadASN1(&inner, cryptobyte_asn1.SEQUENCE) ||
		inner.Empty() ||
		!inner.ReadASN1(&hBytes, cryptobyte_asn1.OCTET_STRING) ||
		!inner.ReadASN1BitString(&sBitString) ||
		!inner.Empty() {
		return data, errors.ErrDecodeASN1Failed
	}
	s.H.SetBytes(hBytes)
	_, err = bn256.UnmarshalG1(&s.S, sBitString.RightAlign())
	if err != nil {
		return data, err
	}
	return []byte(input), err
}

////////////////////////////////////////////////////////////////
//
//						密文结构
//
////////////////////////////////////////////////////////////////

// SM9Cipher ::= SEQUENCE{
// 	EnType	INTEGER,	——加密方式
// 	C1	BIT STRING,	——密文第一部分C1（见GB/T 38635.2—2020）
// 	C3	OCTET STRING,	——密文杂凑值
// 	C2	OCTET STRING 	——密文
// }

func (c *Cipher) MarshalASN1(compressed bool) ([]byte, error) {
	var b cryptobyte.Builder
	c2Bytes := make([]byte, len(c.C)+len(c.IV))
	n := 0
	if c.EncType > 1 {
		n += copy(c2Bytes, c.IV[:])
	}
	n += copy(c2Bytes[n:], c.C)
	c2Bytes = c2Bytes[:n]

	b.AddASN1(cryptobyte_asn1.SEQUENCE, func(b *cryptobyte.Builder) {
		b.AddASN1Int64(int64(c.EncType))
		b.AddASN1BitString(bn256.MarshalG1(&c.C1, compressed))
		b.AddASN1OctetString(c.H[:])
		b.AddASN1OctetString(c2Bytes)
	})
	return b.Bytes()
}

func (c *Cipher) UnmarshalASN1(data []byte) (rest []byte, err error) {
	input := cryptobyte.String(data)
	var inner, c3Bytes, c2Bytes cryptobyte.String
	var c1BitString encoding_asn1.BitString
	var encType uint32
	if !input.ReadASN1(&inner, cryptobyte_asn1.SEQUENCE) ||
		!inner.ReadASN1Integer(&encType) ||
		!inner.ReadASN1BitString(&c1BitString) ||
		!inner.ReadASN1(&c3Bytes, cryptobyte_asn1.OCTET_STRING) ||
		!inner.ReadASN1(&c2Bytes, cryptobyte_asn1.OCTET_STRING) ||
		!inner.Empty() {
		return data, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse Cipher ASN.1 data failed")
	}
	c.EncType = EncType(encType)
	_, err = bn256.UnmarshalG1(&c.C1, c1BitString.RightAlign())
	if err != nil {
		return data, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse Cipher ASN.1 data failed")
	}
	if len(c3Bytes) != len(c.H) {
		return data, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse Cipher ASN.1 data failed")
	}
	copy(c.H[:], c3Bytes)

	if c.EncType > 1 {
		if len(c2Bytes) < len(c.IV) {
			return data, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse Cipher ASN.1 data failed")
		}
		copy(c.IV[:], c2Bytes)
		c2Bytes = c2Bytes[len(c.IV):]
	}
	c.C = append(c.C[:0], c2Bytes...)
	return []byte(input), nil
}

////////////////////////////////////////////////////////////////
//
//						密文封装结构
//
////////////////////////////////////////////////////////////////

// SM9KeyPackage ::= SEQUENCE{
// 	K	OCTET STRING,	——生成的密钥
// 	C	BIT STRING	——封装的交换密文
// }
// K作为用户A保留的密钥。C作为交换密文传递给B用户，B用户利用C可以生成K。

// 这里与标准不同，只是封装C. 密钥不进行封装。即
//
//	SM9KeyPackage := OCTET STRING
func (k *KeyPackage) MarshalASN1(compressed bool) ([]byte, error) {
	var b cryptobyte.Builder
	b.AddASN1BitString(bn256.MarshalG1(&k.G1, compressed))
	return b.Bytes()
}

func (k *KeyPackage) UnmarshalASN1(data []byte) (rest []byte, err error) {
	input := cryptobyte.String(data)
	var b encoding_asn1.BitString
	if ok := input.ReadASN1BitString(&b); !ok {
		return data, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse KeyPackage ASN.1 data failed")
	}
	_, err = bn256.UnmarshalG1(&k.G1, b.RightAlign())
	if err != nil {
		return data, gerrors.WithAnnotating(
			gerrors.ChainErrors(errors.ErrDecodeASN1Failed, err),
			"parse KeyPackage ASN.1 data failed",
		)
	}
	return []byte(input), nil
}