package sdf

import (
	"crypto/cipher"
	"fmt"
	"math/big"
	"strconv"
	"strings"
	"time"

	"xdx.jelly/xgcl/pbkd"
	"xdx.jelly/xgcl/sm/sm2"
	"xdx.jelly/xgcl/sm/sm4"
)

/*
Example of snapshot:
{
	"keys":[
		{
			"type": "Sm2SignKey",
			"url": "sm2sign/12",
			"material": ["base64(d)", "base64(04||x||y)"]
		},
		{
			"type": "Sm2SignKey",
			"url": "sm2enc/12",
			"material": ["base64(d)", "base64(04||x||y)"]
		},
		{
			"type": "BlockKey",
			"url": "kek/0",
			"material": ["base64(key)"]
		}
	]
}

*/
type KeyMaterial struct {
	Type      string   `json:"type"`
	Url       string   `json:"url"`
	Nonces    [][]byte `json:"nonces"` // coresponding to each Material, with add Url
	Materials [][]byte `json:"materials"`
}

type EncParam struct {
	Algo  string `json:"algo"` // only "sm4_gcm" are supported
	Salt  []byte `json:"salt"`
	Count uint32 `json:"count"`
}

type Snapshot struct {
	Time     time.Time      `json:"time"`
	Info     string         `json:"info"`
	Enc      bool           `json:"enc"`
	EncParam *EncParam      `json:"encParam"`
	Keys     []*KeyMaterial `json:"keys"`
}

func (s *SdfNoLock) Snapshot(enc bool, password []byte, count int) (*Snapshot, error) {

	snapshot := &Snapshot{
		Time: time.Now(),
		Enc:  enc,
		Keys: make([]*KeyMaterial, 0),
	}
	var aead cipher.AEAD
	var nonces [][]byte

	if enc {
		var salt []byte
		var err error
		var n uint32
		salt = make([]byte, 16)
		if n, err = s.SDF_GenerateRandom(salt); err != nil || n != uint32(len(salt)) {
			return nil, err
		}
		var encryptionKey []byte
		if encryptionKey, err = pbkd.PbkdfWithHmacSm3(password, salt, count, sm4.BlockSize); err != nil {
			return nil, err
		}

		if aead, err = sm4.NewGCM(encryptionKey); err != nil {
			return nil, err
		}

		nonces = [][]byte{
			make([]byte, aead.NonceSize()),
			make([]byte, aead.NonceSize()),
		}

		snapshot.EncParam = &EncParam{
			Algo:  "sm4_gcm",
			Salt:  salt,
			Count: uint32(count),
		}
	}

	for k, v := range s.Sm2Key {
		if v.signKey != nil {
			sk := v.signKey.Bytes()
			pk := v.signKey.PublicKey.Bytes()
			url := fmt.Sprintf("sm2sign/%d", k)
			thisNonces := nonces
			if aead != nil {
				s.SDF_GenerateRandom(nonces[0])
				s.SDF_GenerateRandom(nonces[1])
				thisNonces = [][]byte{
					append([]byte{}, nonces[0]...),
					append([]byte{}, nonces[1]...),
				}
				sk = aead.Seal(sk[:0], nonces[0], sk, []byte(url))
				pk = aead.Seal(pk[:0], nonces[1], pk, []byte(url))
			}
			km := &KeyMaterial{
				Type:      KeyTypeSm2Sign.String(),
				Url:       url,
				Nonces:    thisNonces,
				Materials: [][]byte{sk, pk},
			}
			snapshot.Keys = append(snapshot.Keys, km)
		}

		if v.encKey != nil {
			sk := v.encKey.Bytes()
			pk := v.encKey.PublicKey.Bytes()
			url := fmt.Sprintf("sm2enc/%d", k)
			thisNonces := nonces
			if aead != nil {
				s.SDF_GenerateRandom(nonces[0])
				s.SDF_GenerateRandom(nonces[1])
				thisNonces = [][]byte{
					append([]byte{}, nonces[0]...),
					append([]byte{}, nonces[1]...),
				}
				sk = aead.Seal(sk[:0], nonces[0], sk, []byte(url))
				pk = aead.Seal(pk[:0], nonces[1], pk, []byte(url))
			}
			km := &KeyMaterial{
				Type:      KeyTypeSm2Enc.String(),
				Url:       url,
				Nonces:    thisNonces,
				Materials: [][]byte{sk, pk},
			}
			snapshot.Keys = append(snapshot.Keys, km)
		}
	}
	for k, v := range s.SymKey {
		key := make([]byte, len(v.key), len(v.key)+16)
		url := fmt.Sprintf("kek/%d", k)
		thisNonces := nonces
		copy(key, v.key)
		if aead != nil {
			s.SDF_GenerateRandom(nonces[0])
			thisNonces = [][]byte{
				append([]byte{}, nonces[0]...),
			}
			key = aead.Seal(key[:0], nonces[0], key, []byte(url))
		}

		km := &KeyMaterial{
			Type:      v.keyType.String(),
			Url:       url,
			Nonces:    thisNonces,
			Materials: [][]byte{key},
		}
		snapshot.Keys = append(snapshot.Keys, km)
	}

	return snapshot, nil
}

func (s *SdfNoLock) Restore(ss *Snapshot, password []byte) error {
	var sdf = &SdfNoLock{
		RsaKey: make(map[uint32]*RsaKey),
		Sm2Key: make(map[uint32]*Sm2Key),
		SymKey: make(map[uint32]*SymKey),
	}
	var aead cipher.AEAD
	var err error
	if ss.Enc {
		var encryptionKey []byte
		if encryptionKey, err = pbkd.PbkdfWithHmacSm3([]byte(password), ss.EncParam.Salt, int(ss.EncParam.Count), sm4.BlockSize); err != nil {
			return err
		}

		if aead, err = sm4.NewGCM(encryptionKey); err != nil {
			return err
		}
	}

	for _, k := range ss.Keys {
		if k == nil {
			return fmt.Errorf("nil key")
		}
		switch GetKeyType(k.Type) {
		case KeyTypeBlock:
			if len(k.Materials) != 1 || len(k.Nonces) != len(k.Materials) {
				return fmt.Errorf("key material error")
			}
			key := k.Materials[0]
			if aead != nil {
				key, err = aead.Open(nil, k.Nonces[0], key, []byte(k.Url))
				if err != nil {
					return err
				}
			}
			urlSplits := strings.Split(k.Url, "/")
			if len(urlSplits) != 2 {
				return fmt.Errorf("key url error")
			}
			if strings.TrimSpace(urlSplits[0]) != "kek" {
				return fmt.Errorf("key url error")
			}
			index, err := strconv.Atoi(urlSplits[1])
			if err != nil {
				return fmt.Errorf("key url error")
			} else {
				sdf.SymKey[uint32(index)] = &SymKey{
					keyType: KeyTypeBlock,
					key:     key,
				}
			}
		case KeyTypeSm2Sign:
			if len(k.Materials) != 2 || len(k.Nonces) != len(k.Materials) {
				return fmt.Errorf("key material error")
			}
			sk := k.Materials[0]
			pk := k.Materials[1]
			if aead != nil {
				sk, err = aead.Open(nil, k.Nonces[0], sk, []byte(k.Url))
				if err != nil {
					return err
				}
				pk, err = aead.Open(nil, k.Nonces[1], pk, []byte(k.Url))
				if err != nil {
					return err
				}
			}
			urlSplits := strings.Split(k.Url, "/")
			if len(urlSplits) != 2 {
				return fmt.Errorf("key url error")
			}
			if strings.TrimSpace(urlSplits[0]) != "sm2sign" {
				return fmt.Errorf("key url error")
			}
			index, err := strconv.Atoi(urlSplits[1])
			if err != nil {
				return fmt.Errorf("key url error")
			} else {
				sm2PrivateKey := &sm2.PrivateKey{
					D: new(big.Int).SetBytes(sk),
				}
				sm2PrivateKey.PublicKey.Curve = sm2.Curve()
				sm2PrivateKey.PublicKey.X = new(big.Int).SetBytes(pk[:sm2.ByteSize()])
				sm2PrivateKey.PublicKey.Y = new(big.Int).SetBytes(pk[sm2.ByteSize():])

				if _, ok := sdf.Sm2Key[uint32(index)]; !ok {
					sdf.Sm2Key[uint32(index)] = &Sm2Key{}
				}
				sdf.Sm2Key[uint32(index)].signKey = sm2PrivateKey
			}

		case KeyTypeSm2Enc:
			if len(k.Materials) != 2 || len(k.Nonces) != len(k.Materials) {
				return fmt.Errorf("key material error")
			}
			sk := k.Materials[0]
			pk := k.Materials[1]
			if aead != nil {
				sk, err = aead.Open(nil, k.Nonces[0], sk, []byte(k.Url))
				if err != nil {
					return err
				}
				pk, err = aead.Open(nil, k.Nonces[1], pk, []byte(k.Url))
				if err != nil {
					return err
				}
			}
			urlSplits := strings.Split(k.Url, "/")
			if len(urlSplits) != 2 {
				return fmt.Errorf("key url error")
			}
			if strings.TrimSpace(urlSplits[0]) != "sm2enc" {
				return fmt.Errorf("key url error")
			}
			index, err := strconv.Atoi(urlSplits[1])
			if err != nil {
				return fmt.Errorf("key url error")
			} else {
				sm2PrivateKey := &sm2.PrivateKey{
					D: new(big.Int).SetBytes(sk),
				}
				sm2PrivateKey.PublicKey.Curve = sm2.Curve()
				sm2PrivateKey.PublicKey.X = new(big.Int).SetBytes(pk[:sm2.ByteSize()])
				sm2PrivateKey.PublicKey.Y = new(big.Int).SetBytes(pk[sm2.ByteSize():])

				if _, ok := sdf.Sm2Key[uint32(index)]; !ok {
					sdf.Sm2Key[uint32(index)] = &Sm2Key{}
				}
				sdf.Sm2Key[uint32(index)].encKey = sm2PrivateKey
			}
		case KeyTypeSm1:
			fallthrough
		case KeyTypeSm4:
			fallthrough
		case KeyTypeSm9MasterSign:
			fallthrough
		case KeyTypeSm9MasterEnc:
			fallthrough
		case KeyTypeSm9UserSign:
			fallthrough
		case KeyTypeSm9UserEnc:
			fallthrough
		case KeyTypeRsa:
			fallthrough
		default:
			return fmt.Errorf("key type errror")
		}
	}

	// now copy sdf to s
	s.Sm2Key = sdf.Sm2Key
	s.SymKey = sdf.SymKey
	s.RsaKey = sdf.RsaKey

	return nil

}