package immessagecrypto

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/binary"
	"hotgo/internal/consts"
	"hotgo/internal/dao"
	"hotgo/utility/encrypt"

	"github.com/gogf/gf/v2/errors/gerror"
)

func GenerateDEK() ([]byte, error) {
	key := make([]byte, encrypt.AESGCMKeySize)
	if _, err := rand.Read(key); err != nil {
		return nil, gerror.Wrap(err, "生成商户消息密钥失败")
	}
	return key, nil
}

func WrapDEK(ctx context.Context, tenantID int64, version uint, dek []byte) (ciphertext, nonce []byte, err error) {
	kek, err := loadRootKEK(ctx)
	if err != nil {
		return nil, nil, err
	}
	return encrypt.AESGCMEncrypt(dek, kek, keyAAD(tenantID, version))
}

func UnwrapDEK(ctx context.Context, tenantID int64, version uint, ciphertext, nonce []byte) ([]byte, error) {
	kek, err := loadRootKEK(ctx)
	if err != nil {
		return nil, err
	}
	dek, err := encrypt.AESGCMDecrypt(ciphertext, nonce, kek, keyAAD(tenantID, version))
	if err != nil {
		return nil, gerror.Wrap(err, "解包商户消息密钥失败")
	}
	if len(dek) != encrypt.AESGCMKeySize {
		return nil, gerror.New("商户消息密钥长度不正确")
	}
	return dek, nil
}

func EncodeDEK(dek []byte) string {
	return base64.StdEncoding.EncodeToString(dek)
}

func loadRootKEK(ctx context.Context) ([]byte, error) {
	columns := dao.SysDictData.Columns()
	value, err := dao.SysDictData.Ctx(ctx).
		Fields(columns.Value).
		Where(columns.Type, consts.ImInternalSecretDictType).
		Where(columns.Label, consts.ImMessageRootKEKLabel).
		Where(columns.Status, consts.StatusEnabled).
		Value()
	if err != nil {
		return nil, gerror.Wrap(err, "读取消息根密钥失败")
	}
	if value.IsEmpty() {
		return nil, gerror.New("消息根密钥未初始化")
	}
	kek, err := base64.StdEncoding.DecodeString(value.String())
	if err != nil || len(kek) != encrypt.AESGCMKeySize {
		return nil, gerror.New("消息根密钥格式不正确")
	}
	return kek, nil
}

func keyAAD(tenantID int64, version uint) []byte {
	aad := make([]byte, 16)
	binary.BigEndian.PutUint64(aad[:8], uint64(tenantID))
	binary.BigEndian.PutUint64(aad[8:], uint64(version))
	return aad
}