package sm9_test

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"math/big"
	"testing"
	"time"

	"xdx.jelly/xgcl/sm/sm9"
)

func decodeHex(h string) []byte {
	buf, err := hex.DecodeString(h)
	if err != nil {
		panic(err)
	}
	return buf
}

func TestIdentifier(t *testing.T) {
	id := sm9.Identifier{
		IdentityType: sm9.OidSm9,
		Alias:        "SM9",
		IdentityData: []byte("alice@xdx.jelly"),
		Serial:       big.NewInt(12345678),
		ValidStart:   time.Now(),
		ValidEnd:     time.Now().Add(time.Hour * 365 * 24),
		IdExtentions: nil,
	}

	b, err := id.MarshalASN1()
	if err != nil {
		t.Fatal(err)
	}

	id2 := &sm9.Identifier{}
	_, err = id2.UnmarshalASN1(b)
	if err != nil {
		t.Fatal(err)
	}
	if !(id.Version == id2.Version) {
		t.Fatal()
	}
	if !(id.IdentityType.Equal(id2.IdentityType)) {
		t.Fatal()
	}
	if !(id.Alias == id2.Alias) {
		t.Fatal()
	}
	if !bytes.Equal(id.IdentityData, id2.IdentityData) {
		t.Fatal()
	}

	// Can't comapre id.ValidStart == id2.ValidStart or
	// id.ValidStart.Equal(id2.ValidStart)
	if !(id.ValidStart.Unix() == id2.ValidStart.Unix()) {
		t.Fatal()
	}
	if !(id.ValidEnd.Unix() == id2.ValidEnd.Unix()) {
		t.Fatal()
	}
	if !bytes.Equal(id.IdExtentions, id2.IdExtentions) {
		t.Fatal()
	}
}

type testKeyVector struct {
	pri               []byte
	priAsn1           []byte
	pubAsn1           []byte
	pubAsn1Compressed []byte
}

var testKeyVectors = map[string]*testKeyVector{
	"MastSignPrivateKey": {
		pri:               decodeHex("000130E78459D78545CB54C587E02CF480CE0B66340F319F348A1D5B1F2DC5F4"),
		priAsn1:           decodeHex("021f0130e78459d78545cb54c587e02cf480ce0b66340f319f348a1d5b1f2dc5f4"),
		pubAsn1:           decodeHex("03818200049f64080b3084f733e48aff4b41b565011ce0711c5e392cfb0ab1b6791b94c40829dba116152d1f786ce843ed24a3b573414d2177386a92dd8f14d65696ea5e3269850938abea0112b57329f447e3a0cbad3e2fdb1a77f335e89e1408d0ef1c2541e00a53dda532da1a7ce027b7a46f741006e85f5cdff0730e75c05fb4e3216d"),
		pubAsn1Compressed: decodeHex("034200039f64080b3084f733e48aff4b41b565011ce0711c5e392cfb0ab1b6791b94c40829dba116152d1f786ce843ed24a3b573414d2177386a92dd8f14d65696ea5e32"),
	},
	"MastEncPrivateKey": {
		pri:               decodeHex("0002E65B0762D042F51F0D23542B13ED8CFA2E9A0E7206361E013A283905E31F"),
		priAsn1:           decodeHex("021f02E65B0762D042F51F0D23542B13ED8CFA2E9A0E7206361E013A283905E31F"),
		pubAsn1:           decodeHex("034200049174542668E8F14AB273C0945C3690C66E5DD09678B86F734C4350567ED0628354E598C6BF749A3DACC9FFFEDD9DB6866C50457CFC7AA2A4AD65C3168FF74210"),
		pubAsn1Compressed: decodeHex("032200029174542668E8F14AB273C0945C3690C66E5DD09678B86F734C4350567ED06283"),
	},
}

func TestASN1MastSignKey(t *testing.T) {
	testVector := testKeyVectors["MastSignPrivateKey"]
	ds := &sm9.MastSignPrivateKey{}
	_, _ = ds.SetBytes(testVector.pri)
	b, err := ds.MarshalASN1()
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(b, testVector.priAsn1) {
		t.Fatal("MastSignPrivateKey marshal failed")
	}
	ds = &sm9.MastSignPrivateKey{}
	_, err = ds.UnmarshalASN1(b)
	if err != nil {
		t.Fatal("MastSignPrivateKey unmarshal error")
	}
	if !bytes.Equal(ds.Bytes(), testVector.pri) {
		fmt.Printf("%x\n", ds.Bytes())

		t.Fatal()
	}

	pubs := ds.Public()
	b, err = pubs.MarshalASN1(false)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(b, testVector.pubAsn1) {
		t.Fatal("MastSignPublicKey marshal failed")
	}
	pubs = &sm9.MastSignPublicKey{}
	_, err = pubs.UnmarshalASN1(b)
	if err != nil {
		t.Fatal("MastSignPublicKey unmarshal error")
	}
	if !bytes.Equal(pubs.Bytes(), ds.Public().Bytes()) {
		t.Fatal()
	}

	b, err = pubs.MarshalASN1(true)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(b, testVector.pubAsn1Compressed) {
		t.Fatal("MastSignPublicKey marshal failed")
	}
	pubs = &sm9.MastSignPublicKey{}
	_, err = pubs.UnmarshalASN1(b)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(pubs.Bytes(), ds.Public().Bytes()) {
		t.Fatal("MastSignPublicKey unmarshal error")
	}
}

func TestASN1MastEncKey(t *testing.T) {
	testVector := testKeyVectors["MastEncPrivateKey"]
	de := &sm9.MastEncPrivateKey{}
	_, _ = de.SetBytes(testVector.pri)
	b, err := de.MarshalASN1()
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(b, testVector.priAsn1) {
		t.Fatal("MastEncPrivateKey marshal failed")
	}
	de = &sm9.MastEncPrivateKey{}
	_, err = de.UnmarshalASN1(b)
	if err != nil {
		t.Fatal("MastEncPrivateKey unmarshal error")
	}
	if !bytes.Equal(de.Bytes(), testVector.pri) {
		t.Fatal()
	}

	pube := de.Public()
	b, err = pube.MarshalASN1(true)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(b, testVector.pubAsn1Compressed) {
		t.Fatal("MastEncPublicKey marshal failed")
	}
	pube = &sm9.MastEncPublicKey{}
	_, err = pube.UnmarshalASN1(b)
	if err != nil {
		t.Fatal("MastEncPublicKey unmarshal error")
	}
	if !bytes.Equal(pube.Bytes(), de.Public().Bytes()) {
		t.Fatal()
	}

	pube = de.Public()
	b, err = pube.MarshalASN1(false)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(b, testVector.pubAsn1) {
		t.Fatal("MastEncPublicKey marshal failed")
	}
	pube = &sm9.MastEncPublicKey{}
	_, err = pube.UnmarshalASN1(b)
	if err != nil {
		t.Fatal("MastEncPublicKey unmarshal error")
	}
	if !bytes.Equal(pube.Bytes(), de.Public().Bytes()) {
		t.Fatal()
	}

}

func TestASN1Signature(t *testing.T) {
	signature := &sm9.Signature{
		H: *new(big.Int).SetBytes(decodeHex("823C4B21E4BD2DFE1ED92C606653E996668563152FC33F55D7BFBB9BD9705ADB")),
	}
	if _, err := signature.S.Unmarshal(decodeHex("73BF96923CE58B6AD0E13E9643A406D8EB98417C50EF1B29CEF9ADB48B6D598C856712F1C2E0968AB7769F42A99586AED139D5B8B3E15891827CC2ACED9BAA05")); err != nil {
		t.Fatal(err)
	}

	// 压缩模式
	b, err := signature.MarshalASN1(true)
	if err != nil {
		t.Fatal(err)
	}
	wantedASN1 := decodeHex("30460420823c4b21e4bd2dfe1ed92c606653e996668563152fc33f55d7bfbb9bd9705adb0322000373bf96923ce58b6ad0e13e9643a406d8eb98417c50ef1b29cef9adb48b6d598c")
	if !bytes.Equal(b, wantedASN1) {
		fmt.Printf("%x\n", b)
		t.Fatal()
	}

	b, err = signature.MarshalASN1(false)
	if err != nil {
		t.Fatal(err)
	}
	wantedASN1 = decodeHex("30660420823C4B21E4BD2DFE1ED92C606653E996668563152FC33F55D7BFBB9BD9705ADB0342000473BF96923CE58B6AD0E13E9643A406D8EB98417C50EF1B29CEF9ADB48B6D598C856712F1C2E0968AB7769F42A99586AED139D5B8B3E15891827CC2ACED9BAA05")
	if !bytes.Equal(b, wantedASN1) {
		fmt.Printf("%x\n", b)
		t.Fatal()
	}
	signature2 := &sm9.Signature{}
	_, err = signature2.UnmarshalASN1(b)
	if err != nil {
		t.Fatal(err)
	}
	if signature.H.Cmp(&signature2.H) != 0 ||
		!signature.S.Equal(&signature2.S) {
		t.Fatal()
	}
}

func TestASN1Cipher(t *testing.T) {
	id := []byte("Bob")
	rnd, _ := hex.DecodeString("0001EDEE3778F441F8DEA3D9FA0ACC4E07EE36C93F9A08618AF4AD85CEDE1C22")
	_, pube, _ := sm9.GenerateMastEncPrivateKey(bytes.NewReader(rnd))

	msg := []byte("Chinese IBE standard")
	// _, _ := sm9.GenerateUserEncKey(id, ke)
	rnd, _ = hex.DecodeString("0000AAC0541779C8FC45E3E2CB25C12B5D2576B2129AE8BB5EE2CBE5EC9E785C")

	c, _ := sm9.Encrypt(sm9.EncTypeKDF, pube, id, msg, bytes.NewReader(rnd), nil)
	var data []byte
	var err error
	if data, err = c.MarshalASN1(false); err != nil ||
		!bytes.Equal(data, decodeHex("307F020100034200042445471164490618E1EE20528FF1D545B0F14C8BCAA44544F03DAB5DAC07D8FF42FFCA97D57CDDC05EA405F2E586FEB3A6930715532B8000759F13059ED59AC00420BA672387BCD6DE5016A158A52BB2E7FC429197BCAB70B25AFEE37A2B9DB9F36704141B5F5B0E951489682F3E64E1378CDD5DA9513B1C")) {
		fmt.Printf("%x\n", data)
		t.Fatal()
	}

	if data, err = c.MarshalASN1(true); err != nil ||
		!bytes.Equal(data, decodeHex("305f020100032200022445471164490618e1ee20528ff1d545b0f14c8bcaa44544f03dab5dac07d8ff0420ba672387bcd6de5016a158a52bb2e7fc429197bcab70b25afee37a2b9db9f36704141b5f5b0e951489682f3e64e1378cdd5da9513b1c")) {
		fmt.Printf("%x\n", data)
		t.Fatal()
	}

	c2 := &sm9.Cipher{}
	rest, err := c2.UnmarshalASN1(data)
	if err != nil || len(rest) > 0 ||
		c.EncType != c2.EncType ||
		c.IV != c2.IV ||
		c.C1 != c2.C1 ||
		c.H != c2.H ||
		!bytes.Equal(c.C, c2.C) {
		t.Fatal()
	}

}