package tokens

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"math/big"
	"strings"
)

// KeyGen is a function that generates a unique key string.
type KeyGen func() (string, error)

// RandomString returns a KeyGen that produces URL-safe base64 random strings.
// The length parameter is the number of random bytes (the encoded output will be longer).
func RandomString(length int) KeyGen {
	return func() (string, error) {
		bytes := make([]byte, length)
		if _, err := rand.Read(bytes); err != nil {
			return "", fmt.Errorf("generating random bytes: %w", err)
		}
		return base64.URLEncoding.EncodeToString(bytes), nil
	}
}

// NumericCode returns a KeyGen that produces a numeric string of the given length.
func NumericCode(length int) KeyGen {
	return func() (string, error) {
		var b strings.Builder
		b.Grow(length)
		for i := 0; i < length; i++ {
			n, err := rand.Int(rand.Reader, big.NewInt(10))
			if err != nil {
				return "", fmt.Errorf("generating random digit: %w", err)
			}
			b.WriteByte(byte('0' + n.Int64()))
		}
		return b.String(), nil
	}
}