package pwhash

import (
	"crypto/subtle"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"strconv"
	"strings"

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

const (
	argonLegacyKeySize  = 32
	argonDefaultKeySize = 16
	argonSaltLen        = 16
)

// Argon2PasswordHash uses the Argon2 hashing algorithm.
type argon2PasswordHash struct {
	// Encoding parameters.
	params argon2Params

	// Codec for string encoding.
	codec argon2Codec
}

// newArgon2PasswordHash returns an Argon2i-based PasswordHash using the
// specified parameters for time, memory, and number of threads.
func newArgon2PasswordHash(kind string, keySize int, time, mem uint32, threads uint8, codec argon2Codec) PasswordHash {
	return &argon2PasswordHash{
		params: argon2Params{
			KeySize: keySize,
			Kind:    kind,
			Time:    time,
			Memory:  mem,
			Threads: threads,
		},
		codec: codec,
	}
}

// NewArgon2 returns an Argon2i-based PasswordHash using the default parameters.
func NewArgon2Legacy() PasswordHash {
	return NewArgon2LegacyWithParams(
		defaultArgon2Params.Time,
		defaultArgon2Params.Memory,
		defaultArgon2Params.Threads,
	)
}

// NewArgon2WithParams returns an Argon2i-based PasswordHash using the
// specified parameters for time, memory, and number of threads.
func NewArgon2LegacyWithParams(time, mem uint32, threads uint8) PasswordHash {
	return newArgon2PasswordHash(kindArgon2I, argonLegacyKeySize, time, mem, threads, &a2LegacyCodec{})
}

// NewArgon2Std returns an Argon2i-based PasswordHash that conforms
// to the reference implementation encoding, using default parameters.
func NewArgon2Std() PasswordHash {
	return NewArgon2StdWithParams(
		defaultArgon2Params.Time,
		defaultArgon2Params.Memory,
		defaultArgon2Params.Threads,
	)
}

// NewArgon2StdWithParams returns an Argon2id-based PasswordHash using
// the specified parameters for time, memory, and number of
// threads. This will use the string encoding ("$argon2id$") documented
// in the argon2 reference implementation.
func NewArgon2StdWithParams(time, mem uint32, threads uint8) PasswordHash {
	return newArgon2PasswordHash(kindArgon2ID, argonDefaultKeySize, time, mem, threads, &argon2StdCodec{})
}

// ComparePassword returns true if the given password matches the
// encrypted one.
func (s *argon2PasswordHash) ComparePassword(encrypted, password string) bool {
	params, salt, dk, err := s.codec.decodeArgon2Hash(encrypted)
	if err != nil {
		return false
	}

	dk2 := params.hash(password, salt)
	return subtle.ConstantTimeCompare(dk, dk2) == 1
}

// Encrypt the given password with the Argon2 algorithm.
func (s *argon2PasswordHash) Encrypt(password string) string {
	salt := getRandomBytes(argonSaltLen)
	dk := s.params.hash(password, salt)
	return s.codec.encodeArgon2Hash(s.params, salt, dk)
}

const (
	kindArgon2I  = "argon2i"
	kindArgon2ID = "argon2id"
)

type argon2Params struct {
	Kind    string
	KeySize int
	Time    uint32
	Memory  uint32
	Threads uint8
}

func (p argon2Params) hash(password string, salt []byte) []byte {
	if p.KeySize == 0 {
		panic("key size is 0")
	}

	switch p.Kind {
	case kindArgon2I:
		return argon2.Key([]byte(password), salt, p.Time, p.Memory, p.Threads, uint32(p.KeySize))
	case kindArgon2ID:
		return argon2.IDKey([]byte(password), salt, p.Time, p.Memory, p.Threads, uint32(p.KeySize))
	default:
		panic("unknown argon2 hash kind")
	}
}

// Default Argon2 parameters are tuned for a high-traffic
// authentication service (<1ms per operation).
var defaultArgon2Params = argon2Params{
	Kind:    kindArgon2ID,
	KeySize: 16,
	Time:    1,
	Memory:  64 * 1024,
	Threads: 4,
}

type argon2Codec interface {
	encodeArgon2Hash(argon2Params, []byte, []byte) string
	decodeArgon2Hash(string) (argon2Params, []byte, []byte, error)
}

// Argon2i legacy encoding, do not use.
type a2LegacyCodec struct{}

func (*a2LegacyCodec) encodeArgon2Hash(params argon2Params, salt, dk []byte) string {
	return fmt.Sprintf("$a2$%d$%d$%d$%x$%x", params.Time, params.Memory, params.Threads, salt, dk)
}

func (*a2LegacyCodec) decodeArgon2Hash(s string) (params argon2Params, salt []byte, dk []byte, err error) {
	if !strings.HasPrefix(s, "$a2$") {
		err = errors.New("not an Argon2 password hash")
		return
	}

	parts := strings.SplitN(s[4:], "$", 5)
	if len(parts) != 5 {
		err = errors.New("bad encoding")
		return
	}

	params.Kind = kindArgon2I

	var i uint64

	if i, err = strconv.ParseUint(parts[0], 10, 32); err != nil {
		return
	}
	params.Time = uint32(i)

	if i, err = strconv.ParseUint(parts[1], 10, 32); err != nil {
		return
	}
	params.Memory = uint32(i)

	if i, err = strconv.ParseUint(parts[2], 10, 8); err != nil {
		return
	}
	params.Threads = uint8(i)

	salt, err = hex.DecodeString(parts[3])
	if err != nil {
		return
	}

	dk, err = hex.DecodeString(parts[4])
	if err != nil {
		return
	}

	params.KeySize = len(dk)
	switch len(dk) {
	case 16, 24, 32:
	default:
		err = errors.New("bad key size")
	}

	return
}

// Standard Argon2 encoding as per the reference implementation in
// https://github.com/P-H-C/phc-winner-argon2/blob/4ac8640c2adc1257677d27d3f833c8d1ee68c7d2/src/encoding.c#L242-L252
type argon2StdCodec struct{}

const argon2HashVersionStr = "v=19"

func (*argon2StdCodec) encodeArgon2Hash(params argon2Params, salt, dk []byte) string {
	encSalt := base64.RawStdEncoding.EncodeToString(salt)
	encDK := base64.RawStdEncoding.EncodeToString(dk)
	return fmt.Sprintf(
		"$%s$%s$m=%d,t=%d,p=%d$%s$%s",
		params.Kind, argon2HashVersionStr,
		params.Memory, params.Time, params.Threads,
		encSalt, encDK)
}

func parseArgon2HashParams(s string) (params argon2Params, err error) {
	params = defaultArgon2Params
	parts := strings.Split(s, ",")
	for _, ss := range parts {
		kv := strings.SplitN(ss, "=", 2)
		if len(kv) != 2 {
			err = errors.New("bad parameter encoding")
			return
		}
		var i uint64
		switch kv[0] {
		case "t":
			i, err = strconv.ParseUint(kv[1], 10, 32)
			params.Time = uint32(i)
		case "m":
			i, err = strconv.ParseUint(kv[1], 10, 32)
			params.Memory = uint32(i)
		case "p":
			i, err = strconv.ParseUint(kv[1], 10, 8)
			params.Threads = uint8(i)
		default:
			err = fmt.Errorf("unknown parameter '%s' in hash", kv[0])
		}
		if err != nil {
			return
		}
	}
	return
}

func (*argon2StdCodec) decodeArgon2Hash(s string) (params argon2Params, salt []byte, dk []byte, err error) {
	var kind string
	switch {
	case strings.HasPrefix(s, "$argon2i$"):
		kind = kindArgon2I
	case strings.HasPrefix(s, "$argon2id$"):
		kind = kindArgon2ID
	default:
		err = errors.New("not an Argon2 password hash")
		return
	}

	parts := strings.SplitN(s, "$", 6)
	if len(parts) != 6 {
		err = errors.New("bad encoding")
		return
	}
	if parts[2] != argon2HashVersionStr {
		err = errors.New("bad argon2 hash version")
		return
	}

	params, err = parseArgon2HashParams(parts[3])
	if err != nil {
		return
	}
	params.Kind = kind

	if salt, err = base64.RawStdEncoding.DecodeString(parts[4]); err != nil {
		return
	}
	if dk, err = base64.RawStdEncoding.DecodeString(parts[5]); err != nil {
		return
	}

	params.KeySize = len(dk)
	switch len(dk) {
	case 16, 24, 32:
	default:
		err = errors.New("bad key size")
	}

	return
}