package security

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"strconv"
	"strings"
)

func pbkdf2SHA256(password, salt []byte, iter, keyLen int) []byte {
	// PBKDF2 per RFC2898
	hLen := 32 // sha256
	numBlocks := (keyLen + hLen - 1) / hLen
	var out []byte
	for block := 1; block <= numBlocks; block++ {
		t := pbkdf2F(password, salt, iter, block)
		out = append(out, t...)
	}
	return out[:keyLen]
}

func pbkdf2F(password, salt []byte, iter, blockNum int) []byte {
	// U1 = PRF(P, S || INT(blockNum))
	// Uc = PRF(P, Uc-1)
	// T = U1 XOR U2 XOR ... XOR Uiter
	b := make([]byte, len(salt)+4)
	copy(b, salt)
	b[len(salt)+0] = byte(blockNum >> 24)
	b[len(salt)+1] = byte(blockNum >> 16)
	b[len(salt)+2] = byte(blockNum >> 8)
	b[len(salt)+3] = byte(blockNum)

	u := hmacSHA256(password, b)
	t := make([]byte, len(u))
	copy(t, u)
	for i := 2; i <= iter; i++ {
		u = hmacSHA256(password, u)
		for j := range t {
			t[j] ^= u[j]
		}
	}
	return t
}

func hmacSHA256(key, msg []byte) []byte {
	m := hmac.New(sha256.New, key)
	m.Write(msg)
	return m.Sum(nil)
}

func HashPasswordPBKDF2(password string, salt []byte, iter int) string {
	key := pbkdf2SHA256([]byte(password), salt, iter, 32)
	return fmt.Sprintf("pbkdf2_sha256$%d$%s$%s",
		iter,
		base64.RawURLEncoding.EncodeToString(salt),
		base64.RawURLEncoding.EncodeToString(key),
	)
}

func VerifyPasswordPBKDF2(password, encoded string) (bool, error) {
	// Go's fmt scanning does not support "scanset" verbs like %[^$]. Parse explicitly.
	parts := strings.Split(encoded, "$")
	if len(parts) != 4 {
		return false, fmt.Errorf("parse hash: expected 4 parts, got %d", len(parts))
	}
	algo := parts[0]
	iter, err := strconv.Atoi(parts[1])
	if err != nil {
		return false, fmt.Errorf("parse hash iter: %w", err)
	}
	saltB64 := parts[2]
	keyB64 := parts[3]
	if algo != "pbkdf2_sha256" {
		return false, fmt.Errorf("unsupported algo %q", algo)
	}
	salt, err := base64.RawURLEncoding.DecodeString(saltB64)
	if err != nil {
		return false, fmt.Errorf("salt decode: %w", err)
	}
	want, err := base64.RawURLEncoding.DecodeString(keyB64)
	if err != nil {
		return false, fmt.Errorf("key decode: %w", err)
	}
	got := pbkdf2SHA256([]byte(password), salt, iter, len(want))
	// constant-time compare
	if len(got) != len(want) {
		return false, nil
	}
	var diff byte
	for i := range got {
		diff |= got[i] ^ want[i]
	}
	return diff == 0, nil
}