package pwhash

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

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

var (
	argonKeyLen  uint32 = 32
	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(time, mem uint32, threads uint8, codec argon2Codec) PasswordHash {
	return &argon2PasswordHash{
		params: argon2Params{
			Time:    time,
			Memory:  mem,
			Threads: threads,
		},
		codec: codec,
	}
}

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

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

// 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 Argon2i-based PasswordHash using
// the specified parameters for time, memory, and number of
// threads. This will use the string encoding ("$argon2$") documented
// in the argon2 reference implementation.
func NewArgon2StdWithParams(time, mem uint32, threads uint8) PasswordHash {
	return newArgon2PasswordHash(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 := argon2.Key([]byte(password), salt, params.Time, params.Memory, params.Threads, argonKeyLen)
	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 := argon2.Key([]byte(password), salt, s.params.Time, s.params.Memory, s.params.Threads, argonKeyLen)
	return s.codec.encodeArgon2Hash(s.params, salt, dk)
}

type argon2Params struct {
	Time    uint32
	Memory  uint32
	Threads uint8
}

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

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

type a2Codec struct{}

func (*a2Codec) 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 (*a2Codec) 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
	}

	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])
	return
}

type argon2StdCodec struct{}

func (*argon2StdCodec) encodeArgon2Hash(params argon2Params, salt, dk []byte) string {
	encSalt := base64.RawStdEncoding.EncodeToString(salt)
	encDK := base64.RawStdEncoding.EncodeToString(dk)
	return fmt.Sprintf("$argon2i$v=19$m=%d,t=%d,p=%d$%s$%s", 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 = errors.New("unknown parameter in hash")
		}
		if err != nil {
			return
		}
	}
	return
}

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

	parts := strings.SplitN(s[9:], "$", 4)
	if len(parts) != 4 {
		err = errors.New("bad encoding")
		return
	}
	if parts[0] != "v=19" {
		err = errors.New("bad argon2 hash version")
		return
	}

	params, err = parseArgon2HashParams(parts[1])
	if err != nil {
		return
	}
	if salt, err = base64.RawStdEncoding.DecodeString(parts[2]); err != nil {
		return
	}
	dk, err = base64.RawStdEncoding.DecodeString(parts[3])

	log.Printf("params: %+v", params)
	return
}