package bn256

import (
	"encoding/asn1"

	"golang.org/x/crypto/cryptobyte"
	"xdx.jelly/xgcl/gerrors"
	"xdx.jelly/xgcl/sm/sm9/errors"
)

// return the bytes 04||x||y or 02(03)||x
func MarshalG1(p *G1, compressed bool) []byte {
	data := make([]byte, 2*numBytes+1)
	// infinity point is encoded as '00'
	if p.IsInfinity() {
		return []byte{0}
	}

	data[0] = 4
	copy(data[1:], p.Marshal())
	if compressed {
		data[0] = 2 + data[2*numBytes]&1
		data = data[:1+numBytes]
	}
	return data
}

func UnmarshalG1(p *G1, b []byte) (rest []byte, err error) {
	// infinity case
	if b[0] == 0 {
		p.SetInfinity()
		return b[1:], nil
	}
	if ((b[0] == 2 || b[0] == 3) && (len(b) < 1+numBytes)) ||
		((b[0] == 4 || b[0] == 6 || b[0] == 7) && (len(b) < 1+2*numBytes)) {
		return b, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse G1 ASN.1 data failed")
	}
	switch b[0] {
	case 2:
		_, _ = p.UnmarshalCompressed(b[1:1+numBytes], 0)
		b = b[1+numBytes:]
	case 3:
		_, _ = p.UnmarshalCompressed(b[1:1+numBytes], 1)
		b = b[1+numBytes:]
	case 6, 7:
		if (b[0] & 1) != b[len(b)-1]&1 {
			return b, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse G1 ASN.1 data failed")
		}
		fallthrough
	case 4:
		_, err := p.Unmarshal(b[1:])
		if err != nil {
			return b, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse G1 ASN.1 data failed")
		}
		b = b[1+2*numBytes:]
	default:
		return b, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse G1 ASN.1 data failed, the leading byte of a EC point must be 2,3,4,6 or 7")
	}
	return b, nil
}

func MarshalG2(p *G2, compressed bool) []byte {
	data := make([]byte, 4*numBytes+1)

	// infinity case, the point is encoded as '00'
	if p.IsInfinity() {
		return []byte{0}
	}

	data[0] = 4
	copy(data[1:], p.Marshal())
	if compressed {
		data[0] = 2 + data[4*numBytes]&1
		data = data[:1+2*numBytes]
	}
	return data
}

func UnmarshalG2(p *G2, b []byte) (rest []byte, err error) {
	// infinity case
	if b[0] == 0 {
		p.SetInfinity()
		return b[1:], nil
	}
	if ((b[0] == 2 || b[0] == 3) && (len(b) < 1+2*numBytes)) ||
		((b[0] == 4 || b[0] == 6 || b[0] == 7) && (len(b) < 1+4*numBytes)) {
		return b, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse G2 ASN.1 data failed")
	}
	switch b[0] {
	case 2:
		_, _ = p.UnmarshalCompressed(b[1:1+numBytes], b[1+numBytes:1+2*numBytes], 0)
		b = b[1+2*numBytes:]
	case 3:
		_, _ = p.UnmarshalCompressed(b[1:1+numBytes], b[1+numBytes:1+2*numBytes], 1)
		b = b[1+2*numBytes:]
	case 6, 7:
		if (b[0] & 1) != b[len(b)-1]&1 {
			return b, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse G2 ASN.1 data failed, the leading byte of a EC point must be 2,3,4,6 or 7")
		}
		fallthrough
	case 4:
		_, err := p.Unmarshal(b[1:])
		if err != nil {
			return b, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse G2 ASN.1 data failed")
		}
		b = b[1+4*numBytes:]
	default:
		return b, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse G2 ASN.1 data failed, the leading byte of a EC point must be 2,3,4,6 or 7")
	}

	return b, nil

}

// //////////////////////////////////////////////////////////////
//
//	ASN.1 编解码G1,G2,
//
// //////////////////////////////////////////////////////////////
func (g *G1) MarshalASN1(compressed bool) ([]byte, error) {
	var b cryptobyte.Builder
	b.AddASN1BitString(MarshalG1(g, compressed))
	return b.Bytes()
}

// UnmarshalASN1 ASN.1解码MastSignPublicKey
//
// 注：支持tag为0,2,3,4,6,7即无穷远点(0)，非压缩形式(4)，压缩形式(2,3)，混合形式(6,7)
func (g *G1) UnmarshalASN1(data []byte) (rest []byte, err error) {
	input := cryptobyte.String(data)
	var b asn1.BitString
	if ok := input.ReadASN1BitString(&b); !ok {
		return data, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse G1 ASN.1 data failed")
	}
	_, err = UnmarshalG1(g, b.RightAlign())
	if err != nil {
		return data, err
	}
	return []byte(input), nil
}

func (g *G2) MarshalASN1(compressed bool) ([]byte, error) {
	var b cryptobyte.Builder
	b.AddASN1BitString(MarshalG2(g, compressed))
	return b.Bytes()
}

// UnmarshalASN1 ASN.1解码MastSignPublicKey
//
// 注：支持tag为0,2,3,4,6,7即无穷远点(0)，非压缩形式(4)，压缩形式(2,3)，混合形式(6,7)
func (g *G2) UnmarshalASN1(data []byte) (rest []byte, err error) {
	input := cryptobyte.String(data)
	var b asn1.BitString
	if ok := input.ReadASN1BitString(&b); !ok {
		return data, gerrors.WithAnnotating(errors.ErrDecodeASN1Failed, "parse G2 ASN.1 data failed")
	}
	_, err = UnmarshalG2(g, b.RightAlign())
	if err != nil {
		return data, err
	}
	return []byte(input), nil
}
