package entities

import (
	"crypto/sha256"
	"errors"
	"go-nkode/config"
	"go-nkode/internal/models"
	"go-nkode/internal/security"
	"golang.org/x/crypto/bcrypt"
)

type UserCipherKeys struct {
	AlphaKey    []uint64
	SetKey      []uint64
	PassKey     []uint64
	MaskKey     []uint64
	Salt        []byte
	MaxNKodeLen int
	Kp          *KeypadDimension
}

func NewUserCipherKeys(kp *KeypadDimension, setVals []uint64, maxNKodeLen int) (*UserCipherKeys, error) {
	err := kp.IsValidKeypadDimension()
	if err != nil {
		return nil, err
	}
	setKey, err := security.GenerateRandomNonRepeatingUint64(kp.AttrsPerKey)
	if err != nil {
		return nil, err
	}
	setKey, err = security.XorLists(setKey, setVals)
	if err != nil {
		return nil, err
	}

	alphaKey, _ := security.GenerateRandomNonRepeatingUint64(kp.TotalAttrs())
	passKey, _ := security.GenerateRandomNonRepeatingUint64(maxNKodeLen)
	maskKey, _ := security.GenerateRandomNonRepeatingUint64(maxNKodeLen)
	salt, _ := security.RandomBytes(10)
	userCipherKeys := UserCipherKeys{
		AlphaKey:    alphaKey,
		PassKey:     passKey,
		MaskKey:     maskKey,
		SetKey:      setKey,
		Salt:        salt,
		MaxNKodeLen: maxNKodeLen,
		Kp:          kp,
	}
	return &userCipherKeys, nil
}

func (u *UserCipherKeys) PadUserMask(userMask []uint64, setVals []uint64) ([]uint64, error) {
	if len(userMask) > u.MaxNKodeLen {
		return nil, config.ErrUserMaskTooLong
	}
	paddedUserMask := make([]uint64, len(userMask))
	copy(paddedUserMask, userMask)
	for i := len(userMask); i < u.MaxNKodeLen; i++ {
		paddedUserMask = append(paddedUserMask, setVals[i%len(setVals)])
	}
	return paddedUserMask, nil
}

func (u *UserCipherKeys) ValidPassword(hashedPassword string, passcodeAttrIdx []int, attrVals []uint64) error {
	hashBytes := []byte(hashedPassword)
	passcodeCipher := u.encipherCode(passcodeAttrIdx, attrVals)
	passwordDigest := u.saltAndDigest(passcodeCipher)
	err := bcrypt.CompareHashAndPassword(hashBytes, passwordDigest)
	if err != nil {
		if errors.Is(err, bcrypt.ErrMismatchedHashAndPassword) {
			return config.ErrInvalidNKode
		}
		return err
	}
	return nil
}

func (u *UserCipherKeys) EncipherSaltHashCode(passcodeAttrIdx []int, attrVals []uint64) (string, error) {
	passcodeCipher := u.encipherCode(passcodeAttrIdx, attrVals)

	passcodeDigest := u.saltAndDigest(passcodeCipher)
	passcodeBytes, err := u.hashPasscode(passcodeDigest)
	if err != nil {
		return "", err
	}
	return string(passcodeBytes), nil
}

func (u *UserCipherKeys) encipherCode(passcodeAttrIdx []int, attrVals []uint64) []uint64 {
	passcodeLen := len(passcodeAttrIdx)

	passcodeCipher := make([]uint64, u.MaxNKodeLen)
	for idx := 0; idx < passcodeLen; idx++ {
		attrIdx := passcodeAttrIdx[idx]
		alpha := u.AlphaKey[attrIdx]
		attrVal := attrVals[attrIdx]
		pass := u.PassKey[idx]
		passcodeCipher[idx] = alpha ^ pass ^ attrVal
	}
	return passcodeCipher
}

func (u *UserCipherKeys) saltAndDigest(passcode []uint64) []byte {
	passcodeBytes := security.Uint64ArrToByteArr(passcode)
	passcodeBytes = append(passcodeBytes, u.Salt...)
	h := sha256.New()
	h.Write(passcodeBytes)

	return h.Sum(nil)
}

func (u *UserCipherKeys) hashPasscode(passcodeDigest []byte) ([]byte, error) {
	hashedPassword, err := bcrypt.GenerateFromPassword(passcodeDigest, bcrypt.DefaultCost)
	if err != nil {
		return nil, err
	}
	return hashedPassword, nil
}

func (u *UserCipherKeys) EncipherMask(passcodeSet []uint64, customerAttrs CustomerAttributes, userKp KeypadDimension) (string, error) {
	setVals, err := customerAttrs.SetValsForKp(userKp)
	if err != nil {
		return "", err
	}
	paddedPasscodeSets, err := u.PadUserMask(passcodeSet, setVals)
	if err != nil {
		return "", err
	}

	cipheredMask := make([]uint64, len(paddedPasscodeSets))
	for idx := range paddedPasscodeSets {
		setIdx, err := customerAttrs.IndexOfSet(paddedPasscodeSets[idx])
		if err != nil {
			return "", err
		}
		setKeyVal := u.SetKey[setIdx]
		maskKeyVal := u.MaskKey[idx]
		setVal := paddedPasscodeSets[idx]
		cipheredMask[idx] = setKeyVal ^ maskKeyVal ^ setVal
	}
	mask := security.EncodeBase64Str(cipheredMask)
	return mask, nil
}

func (u *UserCipherKeys) DecipherMask(mask string, setVals []uint64, passcodeLen int) ([]uint64, error) {
	decodedMask, err := security.DecodeBase64Str(mask)
	if err != nil {
		return nil, err
	}
	decipheredMask, err := security.XorLists(decodedMask, u.MaskKey)
	if err != nil {
		return nil, err
	}
	setKeyRandComponent, err := security.XorLists(setVals, u.SetKey)
	if err != nil {
		return nil, err
	}

	passcodeSet := make([]uint64, passcodeLen)
	for idx, setCipher := range decipheredMask[:passcodeLen] {
		setIdx, err := security.IndexOf(setKeyRandComponent, setCipher)
		if err != nil {
			return nil, err
		}
		passcodeSet[idx] = setVals[setIdx]
	}
	return passcodeSet, nil
}

func (u *UserCipherKeys) EncipherNKode(passcodeAttrIdx []int, customerAttrs CustomerAttributes) (*models.EncipheredNKode, error) {
	attrVals, err := customerAttrs.AttrValsForKp(*u.Kp)
	code, err := u.EncipherSaltHashCode(passcodeAttrIdx, attrVals)
	if err != nil {
		return nil, err
	}
	passcodeSet := make([]uint64, len(passcodeAttrIdx))

	for idx := range passcodeSet {
		passcodeAttr := attrVals[passcodeAttrIdx[idx]]
		passcodeSet[idx], err = customerAttrs.GetAttrSetVal(passcodeAttr, *u.Kp)
		if err != nil {
			return nil, err
		}
	}
	mask, err := u.EncipherMask(passcodeSet, customerAttrs, *u.Kp)
	if err != nil {
		return nil, err
	}
	encipheredCode := models.EncipheredNKode{
		Code: code,
		Mask: mask,
	}
	return &encipheredCode, nil
}
