package identity

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"strconv"
	"strings"
)

const (
	currentPBKDF2Iterations = 600_000
	legacyPBKDF2Iterations  = 120_000
	maximumPBKDF2Iterations = 2_000_000
	passwordDigestBytes     = 32
)

var ErrInvalidPasswordHash = errors.New("invalid password hash")

type PasswordHasher struct{}

func (PasswordHasher) Hash(password string) (string, error) {
	saltBytes := make([]byte, 16)
	if _, err := rand.Read(saltBytes); err != nil {
		return "", fmt.Errorf("generate password salt: %w", err)
	}
	return hashWithSalt(password, hex.EncodeToString(saltBytes), currentPBKDF2Iterations), nil
}

func (PasswordHasher) Verify(password, stored string) bool {
	iterations, salt, expected, err := parsePasswordHash(stored)
	if err != nil {
		return false
	}
	actual := pbkdf2SHA256([]byte(password), []byte(salt), iterations, len(expected))
	return subtle.ConstantTimeCompare(actual, expected) == 1
}

func (PasswordHasher) NeedsUpgrade(stored string) bool {
	iterations, _, _, err := parsePasswordHash(stored)
	return err == nil && iterations < currentPBKDF2Iterations
}

func hashWithSalt(password, salt string, iterations int) string {
	digest := pbkdf2SHA256([]byte(password), []byte(salt), iterations, passwordDigestBytes)
	return fmt.Sprintf("pbkdf2_sha256$%d$%s$%s", iterations, salt, hex.EncodeToString(digest))
}

func parsePasswordHash(stored string) (int, string, []byte, error) {
	parts := strings.Split(stored, "$")
	iterations := legacyPBKDF2Iterations
	var salt, encodedDigest string
	switch {
	case len(parts) == 4 && parts[0] == "pbkdf2_sha256":
		parsed, err := strconv.Atoi(parts[1])
		if err != nil {
			return 0, "", nil, ErrInvalidPasswordHash
		}
		iterations, salt, encodedDigest = parsed, parts[2], parts[3]
	case len(parts) == 2:
		salt, encodedDigest = parts[0], parts[1]
	default:
		return 0, "", nil, ErrInvalidPasswordHash
	}
	if iterations < 1 || iterations > maximumPBKDF2Iterations || salt == "" {
		return 0, "", nil, ErrInvalidPasswordHash
	}
	digest, err := hex.DecodeString(encodedDigest)
	if err != nil || len(digest) != passwordDigestBytes {
		return 0, "", nil, ErrInvalidPasswordHash
	}
	return iterations, salt, digest, nil
}

func pbkdf2SHA256(password, salt []byte, iterations, keyLength int) []byte {
	const hashLength = sha256.Size
	blocks := (keyLength + hashLength - 1) / hashLength
	result := make([]byte, 0, blocks*hashLength)
	buffer := make([]byte, len(salt)+4)
	copy(buffer, salt)
	for block := 1; block <= blocks; block++ {
		binary.BigEndian.PutUint32(buffer[len(salt):], uint32(block))
		mac := hmac.New(sha256.New, password)
		_, _ = mac.Write(buffer)
		u := mac.Sum(nil)
		t := append([]byte(nil), u...)
		for round := 1; round < iterations; round++ {
			mac.Reset()
			_, _ = mac.Write(u)
			u = mac.Sum(nil)
			for index := range t {
				t[index] ^= u[index]
			}
		}
		result = append(result, t...)
	}
	return result[:keyLength]
}