///
/// Copyright (c) 2018 xdx. All rights reserved.
///
/// \file: enc_impl.go
///
/// \brief: sm9加密辅助函数
///
/// \author: xdx
///

package sm9

import (
	"bytes"
	"io"
	"math/big"

	"xdx.jelly/xgcl/gerrors"
	"xdx.jelly/xgcl/sm/sm4"
	"xdx.jelly/xgcl/sm/sm9/errors"
)

type encReader struct {
	E encoder
	R io.Reader
}

func (er *encReader) Read(dst []byte) (n int, err error) {
	n, err = er.R.Read(dst)
	if err != nil {
		return
	}
	dst, err = er.E.encode(dst, dst)
	return len(dst), nil
}

type encoder interface {
	encode(dst, src []byte) ([]byte, error)
}

type encFunc func(dst, iv, key, src []byte) ([]byte, error)
type worker struct {
	key []byte
	iv  []byte
	f   encFunc
}

func (o *worker) encode(dst, src []byte) ([]byte, error) {
	return o.f(dst, o.iv, o.key, src)
}

func DecryptParams(c *Cipher) (int, func(dst, iv, key, src []byte) ([]byte, error), error) {
	switch c.EncType {
	case EncTypeKDF:
		return len(c.C) + macKeySize, xor, nil
	case EncTypeSM4ECB:
		return sm4.BlockSize + macKeySize, decEcb, nil
	case EncTypeSM4CBC:
		return sm4.BlockSize + macKeySize, decCbc, nil
	case EncTypeSM4CFB:
		return sm4.BlockSize + macKeySize, decCfb, nil
	case EncTypeSM4OFB:
		return sm4.BlockSize + macKeySize, decOfb, nil
	default:
		return 0, nil, errors.ErrEncUnsupportedMode
	}
}

var DecryptFunc = map[EncType]func(dst, iv, key, src []byte) ([]byte, error){
	EncTypeKDF:    xor,
	EncTypeSM4ECB: decEcb,
	EncTypeSM4CBC: decCbc,
	EncTypeSM4CFB: decCfb,
	EncTypeSM4OFB: decOfb,
}

func xor(dst, iv, key, src []byte) ([]byte, error) {
	for i := range src {
		dst[i] = key[i] ^ src[i]
	}
	return dst, nil
}

func encEcb(dst, iv, key, src []byte) ([]byte, error) {
	return sm4.ECBEncrypt(dst, key, src)
}

func decEcb(dst, iv, key, src []byte) ([]byte, error) {
	return sm4.ECBDecrypt(dst, key, src)
}

func encCbc(dst, iv, key, src []byte) ([]byte, error) {
	return sm4.CBCEncrypt(dst, iv, key, src)
}
func decCbc(dst, iv, key, src []byte) ([]byte, error) {
	return sm4.CBCDecrypt(dst, iv, key, src)
}

func encCfb(dst, iv, key, src []byte) ([]byte, error) {
	return sm4.CFBEncrypt(dst, iv, key, src)
}

func decCfb(dst, iv, key, src []byte) ([]byte, error) {
	return sm4.CFBDecrypt(dst, iv, key, src)
}

func encOfb(dst, iv, key, src []byte) ([]byte, error) {
	return sm4.OFBEncrypt(dst, iv, key, src)
}

func decOfb(dst, iv, key, src []byte) ([]byte, error) {
	return sm4.OFBDecrypt(dst, iv, key, src)
}

// assume plain is padding and n blocks, iv is set of 16 bytes if cbc
func encryption(id, plain []byte, pube *MastEncPublicKey, iv, rnd []byte, keylen int, f encFunc, encType EncType) *Cipher {
	c := new(Cipher)
	c.C = make([]byte, len(plain))
	r := new(big.Int).SetBytes(rnd)
	// c.X, c.Y = genKeyPackage(id, pube, rnd)
	c.C1.Set(hashToG1(id, &pube.G1, hidEncryption))
	c.C1.ScalarMult(&c.C1, r)
	c.EncType = encType
	copy(c.IV[:], iv)

	key := make([]byte, keylen)

	// genKey(key, id, &c.C1, &pube.G1, g2Gen, r)
	w := &GT{}
	if pube.e != nil {
		w.ScalarMult(pube.e, r)
	} else {
		pairing(w, &pube.G1, g2Gen)
		w.ScalarMult(w, r)
	}
	genKey(key, id, &c.C1, w)

	reader := &encReader{
		E: &worker{key[:len(key)-macKeySize], iv, f},
		R: bytes.NewReader(plain),
	}
	_, _ = reader.Read(c.C)
	copy(c.H[:], mac(key[len(key)-macKeySize:], c.C))
	return c
}

// decryptions, encType are not used
func decryption(id []byte, c *Cipher, de *UserEncKey, keylen int, f encFunc) ([]byte, error) {
	if !c.C1.IsValid() {
		return nil, gerrors.WithAnnotating(errors.ErrDecFailed, "C1 is not a valid point on curve")
	}

	w := &GT{}
	pairing(w, &c.C1, &de.G2)

	key := make([]byte, keylen)
	genKey(key, id, &c.C1, w)
	return decodeCipher(c, key, f)
}

// DecodeCipher export decodeCipher for tpc package.
// Do not use this function unless you know what you are doing.
var DecodeCipher = decodeCipher

func decodeCipher(c *Cipher, key []byte, f encFunc) ([]byte, error) {
	keylen := len(key)

	m := mac(key[keylen-macKeySize:], c.C)
	if !bytes.Equal(c.H[:], m[:]) {
		return nil, gerrors.WithAnnotating(errors.ErrDecFailed, "check C3 failed")
	}

	reader := &encReader{
		E: &worker{key[:keylen-macKeySize], c.IV[:], f},
		R: bytes.NewReader(c.C),
	}
	plain := make([]byte, len(c.C))
	_, _ = reader.Read(plain)

	return plain, nil
}