package password

import (
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"errors"
	"fmt"
	"sort"
	"strings"

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

const (
	AlgorithmArgon2id        = "argon2id"
	AlgorithmPBKDF2SHA256    = "pbkdf2-hmac-sha256"
	AlgorithmPrototypeSHA256 = "prototype-sha256"
)

var ErrUnsupportedKDF = errors.New("unsupported password KDF")

type Record struct {
	Algorithm string `json:"algorithm"`
	Version   int    `json:"version"`
	Salt      string `json:"salt"`
	Hash      string `json:"hash"`
	Params    Params `json:"params"`
}

type Params struct {
	MemoryKiB   uint32 `json:"memory_kib,omitempty"`
	Iterations  uint32 `json:"iterations,omitempty"`
	Parallelism uint8  `json:"parallelism,omitempty"`
	KeyLength   uint32 `json:"key_length,omitempty"`
}

type KDF interface {
	Algorithm() string
	Hash(password, salt string) (Record, error)
	Verify(record Record, password string) (bool, error)
}

type Registry struct {
	kdfs map[string]KDF
}

func NewRegistry(kdfs ...KDF) Registry {
	registry := Registry{kdfs: make(map[string]KDF, len(kdfs))}
	for _, kdf := range kdfs {
		if kdf != nil {
			registry.kdfs[kdf.Algorithm()] = kdf
		}
	}
	return registry
}

func NewPrototypeRegistry() Registry {
	return NewRegistry(PrototypeSHA256{})
}

func NewProductionRegistry() Registry {
	return NewRegistry(NewArgon2id(Argon2idDefaultParams()), PrototypeSHA256{})
}

func (r Registry) Verify(record Record, password string) (bool, error) {
	kdf, ok := r.kdfs[record.Algorithm]
	if !ok {
		return false, fmt.Errorf("%w: %s", ErrUnsupportedKDF, record.Algorithm)
	}
	return kdf.Verify(record, password)
}

func (r Registry) Supports(algorithm string) bool {
	_, ok := r.kdfs[algorithm]
	return ok
}

func (r Registry) Algorithms() []string {
	algorithms := make([]string, 0, len(r.kdfs))
	for algorithm := range r.kdfs {
		algorithms = append(algorithms, algorithm)
	}
	sort.Strings(algorithms)
	return algorithms
}

type PrototypeSHA256 struct{}

func (PrototypeSHA256) Algorithm() string {
	return AlgorithmPrototypeSHA256
}

func (PrototypeSHA256) DefaultAuthSeedCacheKey() string {
	return AlgorithmPrototypeSHA256 + ":v1"
}

func (PrototypeSHA256) Hash(passwordValue, salt string) (Record, error) {
	if strings.TrimSpace(salt) == "" {
		return Record{}, errors.New("password salt must not be empty")
	}
	return Record{
		Algorithm: AlgorithmPrototypeSHA256,
		Version:   1,
		Salt:      salt,
		Hash:      prototypeHash(salt, passwordValue),
		Params: Params{
			KeyLength: 32,
		},
	}, nil
}

func (kdf PrototypeSHA256) Verify(record Record, passwordValue string) (bool, error) {
	if record.Algorithm != kdf.Algorithm() {
		return false, fmt.Errorf("%w: %s", ErrUnsupportedKDF, record.Algorithm)
	}
	if record.Version != 1 || strings.TrimSpace(record.Salt) == "" || record.Hash == "" {
		return false, errors.New("invalid prototype password hash record")
	}
	next, err := kdf.Hash(passwordValue, record.Salt)
	if err != nil {
		return false, err
	}
	return next.Hash == record.Hash, nil
}

func Argon2idPolicyRecord() Record {
	return Record{
		Algorithm: AlgorithmArgon2id,
		Version:   1,
		Params:    Argon2idDefaultParams(),
	}
}

type Argon2id struct {
	params Params
}

func NewArgon2id(params Params) Argon2id {
	return Argon2id{params: normalizeArgon2idParams(params)}
}

func Argon2idDefaultParams() Params {
	return Params{
		MemoryKiB:   64 * 1024,
		Iterations:  3,
		Parallelism: 1,
		KeyLength:   32,
	}
}

func (Argon2id) Algorithm() string {
	return AlgorithmArgon2id
}

func (kdf Argon2id) DefaultAuthSeedCacheKey() string {
	params := normalizeArgon2idParams(kdf.params)
	return fmt.Sprintf(
		"%s:v1:m=%d:t=%d:p=%d:l=%d",
		AlgorithmArgon2id,
		params.MemoryKiB,
		params.Iterations,
		params.Parallelism,
		params.KeyLength,
	)
}

func (kdf Argon2id) Hash(passwordValue, salt string) (Record, error) {
	if strings.TrimSpace(salt) == "" {
		return Record{}, errors.New("password salt must not be empty")
	}
	params := normalizeArgon2idParams(kdf.params)
	hash := argon2.IDKey(
		[]byte(passwordValue),
		[]byte(salt),
		params.Iterations,
		params.MemoryKiB,
		params.Parallelism,
		params.KeyLength,
	)
	return Record{
		Algorithm: AlgorithmArgon2id,
		Version:   1,
		Salt:      salt,
		Hash:      base64.StdEncoding.EncodeToString(hash),
		Params:    params,
	}, nil
}

func (kdf Argon2id) Verify(record Record, passwordValue string) (bool, error) {
	if record.Algorithm != kdf.Algorithm() {
		return false, fmt.Errorf("%w: %s", ErrUnsupportedKDF, record.Algorithm)
	}
	if record.Version != 1 || strings.TrimSpace(record.Salt) == "" || record.Hash == "" {
		return false, errors.New("invalid argon2id password hash record")
	}
	expected, err := argon2idKDFFromRecord(record).Hash(passwordValue, record.Salt)
	if err != nil {
		return false, err
	}
	return subtle.ConstantTimeCompare([]byte(expected.Hash), []byte(record.Hash)) == 1, nil
}

func argon2idKDFFromRecord(record Record) Argon2id {
	return NewArgon2id(record.Params)
}

func normalizeArgon2idParams(params Params) Params {
	defaults := Argon2idDefaultParams()
	if params.MemoryKiB == 0 {
		params.MemoryKiB = defaults.MemoryKiB
	}
	if params.Iterations == 0 {
		params.Iterations = defaults.Iterations
	}
	if params.Parallelism == 0 {
		params.Parallelism = defaults.Parallelism
	}
	if params.KeyLength == 0 {
		params.KeyLength = defaults.KeyLength
	}
	return params
}

func NewSalt() (string, error) {
	var raw [16]byte
	if _, err := rand.Read(raw[:]); err != nil {
		return "", fmt.Errorf("generate password salt: %w", err)
	}
	return base64.RawURLEncoding.EncodeToString(raw[:]), nil
}

func PBKDF2SHA256PolicyRecord() Record {
	return Record{
		Algorithm: AlgorithmPBKDF2SHA256,
		Version:   1,
		Params:    PBKDF2SHA256DefaultParams(),
	}
}

type PBKDF2SHA256 struct {
	params Params
}

func NewPBKDF2SHA256(params Params) PBKDF2SHA256 {
	return PBKDF2SHA256{params: normalizePBKDF2SHA256Params(params)}
}

func PBKDF2SHA256DefaultParams() Params {
	return Params{
		Iterations: 310000,
		KeyLength:  32,
	}
}

func (PBKDF2SHA256) Algorithm() string {
	return AlgorithmPBKDF2SHA256
}

func (kdf PBKDF2SHA256) DefaultAuthSeedCacheKey() string {
	params := normalizePBKDF2SHA256Params(kdf.params)
	return fmt.Sprintf(
		"%s:v1:i=%d:l=%d",
		AlgorithmPBKDF2SHA256,
		params.Iterations,
		params.KeyLength,
	)
}

func (kdf PBKDF2SHA256) Hash(passwordValue, salt string) (Record, error) {
	if strings.TrimSpace(salt) == "" {
		return Record{}, errors.New("password salt must not be empty")
	}
	params := normalizePBKDF2SHA256Params(kdf.params)
	hash := pbkdf2.Key([]byte(passwordValue), []byte(salt), int(params.Iterations), int(params.KeyLength), sha256.New)
	return Record{
		Algorithm: AlgorithmPBKDF2SHA256,
		Version:   1,
		Salt:      salt,
		Hash:      base64.StdEncoding.EncodeToString(hash),
		Params:    params,
	}, nil
}

func (kdf PBKDF2SHA256) Verify(record Record, passwordValue string) (bool, error) {
	if record.Algorithm != kdf.Algorithm() {
		return false, fmt.Errorf("%w: %s", ErrUnsupportedKDF, record.Algorithm)
	}
	if record.Version != 1 || strings.TrimSpace(record.Salt) == "" || record.Hash == "" {
		return false, errors.New("invalid pbkdf2 password hash record")
	}
	expected, err := pbkdf2SHA256KDFFromRecord(record).Hash(passwordValue, record.Salt)
	if err != nil {
		return false, err
	}
	return subtle.ConstantTimeCompare([]byte(expected.Hash), []byte(record.Hash)) == 1, nil
}

func pbkdf2SHA256KDFFromRecord(record Record) PBKDF2SHA256 {
	return NewPBKDF2SHA256(record.Params)
}

func normalizePBKDF2SHA256Params(params Params) Params {
	defaults := PBKDF2SHA256DefaultParams()
	if params.Iterations == 0 {
		params.Iterations = defaults.Iterations
	}
	if params.KeyLength == 0 {
		params.KeyLength = defaults.KeyLength
	}
	return Params{
		Iterations: params.Iterations,
		KeyLength:  params.KeyLength,
	}
}

func prototypeHash(salt, passwordValue string) string {
	sum := sha256.Sum256([]byte(salt + ":" + passwordValue))
	return base64.StdEncoding.EncodeToString(sum[:])
}