package hashx

import (
	"crypto/pbkdf2"
	"crypto/sha256"
	"crypto/sha512"
	"errors"

	"golang.org/x/crypto/argon2"
)

var (
	ErrInvalidKDFSalt       = errors.New("kdf salt must be non-empty")
	ErrInvalidKDFIterations = errors.New("kdf iterations must be > 0")
	ErrInvalidKDFKeyLength  = errors.New("kdf key length must be > 0")
	ErrInvalidArgon2Params  = errors.New("argon2 params must have time, memory, threads, and key length > 0")
)

// Argon2Params configures Argon2 key derivation.
type Argon2Params struct {
	Time    uint32
	Memory  uint32
	Threads uint8
	KeyLen  uint32
}

// DefaultArgon2idParams returns a conservative default suitable for general online usage.
func DefaultArgon2idParams() Argon2Params {
	return Argon2Params{
		Time:    1,
		Memory:  64 * 1024, // 64 MiB in KiB
		Threads: 4,
		KeyLen:  32,
	}
}

func validatePBKDF2Params(salt []byte, iterations, keyLen int) error {
	if len(salt) == 0 {
		return ErrInvalidKDFSalt
	}
	if iterations <= 0 {
		return ErrInvalidKDFIterations
	}
	if keyLen <= 0 {
		return ErrInvalidKDFKeyLength
	}
	return nil
}

func validateArgon2Params(salt []byte, params Argon2Params) error {
	if len(salt) == 0 {
		return ErrInvalidKDFSalt
	}
	if params.Time == 0 || params.Memory == 0 || params.Threads == 0 || params.KeyLen == 0 {
		return ErrInvalidArgon2Params
	}
	return nil
}

// DerivePBKDF2SHA256Key derives a key with PBKDF2-HMAC-SHA256.
func DerivePBKDF2SHA256Key(password string, salt []byte, iterations, keyLen int) ([]byte, error) {
	if err := validatePBKDF2Params(salt, iterations, keyLen); err != nil {
		return nil, err
	}
	return pbkdf2.Key(sha256.New, password, salt, iterations, keyLen)
}

// DerivePBKDF2SHA512Key derives a key with PBKDF2-HMAC-SHA512.
func DerivePBKDF2SHA512Key(password string, salt []byte, iterations, keyLen int) ([]byte, error) {
	if err := validatePBKDF2Params(salt, iterations, keyLen); err != nil {
		return nil, err
	}
	return pbkdf2.Key(sha512.New, password, salt, iterations, keyLen)
}

// DeriveArgon2idKey derives a key with Argon2id.
func DeriveArgon2idKey(password string, salt []byte, params Argon2Params) ([]byte, error) {
	if err := validateArgon2Params(salt, params); err != nil {
		return nil, err
	}
	return argon2.IDKey([]byte(password), salt, params.Time, params.Memory, params.Threads, params.KeyLen), nil
}

// DeriveArgon2iKey derives a key with Argon2i.
func DeriveArgon2iKey(password string, salt []byte, params Argon2Params) ([]byte, error) {
	if err := validateArgon2Params(salt, params); err != nil {
		return nil, err
	}
	return argon2.Key([]byte(password), salt, params.Time, params.Memory, params.Threads, params.KeyLen), nil
}