package userapi

import (
	"context"
	"crypto/rand"
	"crypto/subtle"
	"errors"
	"fmt"

	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// Registration token errors.
var (
	ErrNoSuchRegistrationToken = storage.ErrNoSuchRegistrationToken
	ErrRegistrationTokenExists = storage.ErrRegistrationTokenExists

	// ErrInvalidRegistrationToken is a token an admin asked for that does
	// not fit the spec's grammar.
	ErrInvalidRegistrationToken = errors.New("invalid registration token")
)

// MaxRegistrationTokenLength is the spec's limit on a registration token.
const MaxRegistrationTokenLength = 64

// DefaultRegistrationTokenLength is what a generated token is when the
// admin does not say.
const DefaultRegistrationTokenLength = 16

// registrationTokenAlphabet is the spec's character set for a token, less
// the punctuation: a generated token is read off a screen and typed.
const registrationTokenAlphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"

// RegistrationTokens answers the m.login.registration_token stage from two
// places: the static tokens of the config file, which never run out, and
// the tokens an admin made through the admin API, which may.
type RegistrationTokens struct {
	store    storage.AdminStore
	static   []string
	required bool
	nowMS    func() int64
}

// RegistrationTokensOptions configure them.
type RegistrationTokensOptions struct {
	// Store holds the tokens made through the admin API. Nil means there
	// are none and the admin endpoints cannot make any.
	Store storage.AdminStore

	// Static are registration.tokens from the config file.
	Static []string

	// Required puts the stage into the registration flow even with no
	// static tokens — registration.require_token, for a server whose
	// tokens all come from the admin API.
	Required bool

	NowMS func() int64
}

// NewRegistrationTokens builds them.
func NewRegistrationTokens(opts RegistrationTokensOptions) *RegistrationTokens {
	return &RegistrationTokens{
		store:    opts.Store,
		static:   opts.Static,
		required: opts.Required || len(opts.Static) > 0,
		nowMS:    opts.NowMS,
	}
}

// Required reports whether registration needs a token.
func (t *RegistrationTokens) Required() bool { return t.required }

// Valid reports whether a token would be accepted now. It does not spend
// it: the stage is checked when it is passed and spent when the account is
// created, so a registration abandoned between the two costs no use.
func (t *RegistrationTokens) Valid(ctx context.Context, token string) (bool, error) {
	if t.staticMatch(token) {
		return true, nil
	}
	if t.store == nil || !wellFormedToken(token) {
		return false, nil
	}
	stored, err := t.store.RegistrationToken(ctx, token)
	switch {
	case errors.Is(err, storage.ErrNoSuchRegistrationToken):
		return false, nil
	case err != nil:
		return false, err
	}
	return stored.Usable(t.nowMS()), nil
}

// Spend takes one use of a token for an account about to be created, and
// returns the function that gives it back if the creation then fails. A
// static token has no uses to count. A token that has run out or expired
// since the stage was passed is ErrNoSuchRegistrationToken.
func (t *RegistrationTokens) Spend(ctx context.Context, token string) (release func(), err error) {
	if t.staticMatch(token) {
		return func() {}, nil
	}
	if t.store == nil {
		return nil, ErrNoSuchRegistrationToken
	}
	if err := t.store.UseRegistrationToken(ctx, token, t.nowMS()); err != nil {
		return nil, err
	}
	return func() {
		// The request's context may be what failed; giving the use back
		// must not depend on it.
		_ = t.store.ReleaseRegistrationToken(context.WithoutCancel(ctx), token)
	}, nil
}

// staticMatch compares a token against every config file token in
// constant time, and does not stop at the first match, so neither the
// value nor the position of a token leaks through the time it takes.
func (t *RegistrationTokens) staticMatch(token string) bool {
	found := 0
	for _, want := range t.static {
		found |= subtle.ConstantTimeCompare([]byte(token), []byte(want))
	}
	return found == 1
}

// NewRegistrationToken is what the admin API asks for.
type NewRegistrationToken struct {
	// Token is the token to make; empty generates one of Length
	// characters.
	Token  string
	Length int

	UsesAllowed *int
	ExpiresAtMS int64
}

// CreateRegistrationToken makes a token through the admin API.
func (t *RegistrationTokens) CreateRegistrationToken(ctx context.Context, req NewRegistrationToken) (storage.RegistrationToken, error) {
	if t.store == nil {
		return storage.RegistrationToken{}, storage.ErrNotImplemented
	}
	token := req.Token
	if token == "" {
		length := req.Length
		if length == 0 {
			length = DefaultRegistrationTokenLength
		}
		if length < 1 || length > MaxRegistrationTokenLength {
			return storage.RegistrationToken{}, fmt.Errorf("%w: length must be between 1 and %d",
				ErrInvalidRegistrationToken, MaxRegistrationTokenLength)
		}
		var err error
		if token, err = generateRegistrationToken(length); err != nil {
			return storage.RegistrationToken{}, err
		}
	} else if !wellFormedToken(token) {
		return storage.RegistrationToken{}, fmt.Errorf(
			"%w: a token is 1 to %d of the characters A-Z, a-z, 0-9, '.', '_', '~' and '-'",
			ErrInvalidRegistrationToken, MaxRegistrationTokenLength)
	}
	if req.UsesAllowed != nil && *req.UsesAllowed < 0 {
		return storage.RegistrationToken{}, fmt.Errorf("%w: uses_allowed must not be negative", ErrInvalidRegistrationToken)
	}
	if req.ExpiresAtMS < 0 {
		return storage.RegistrationToken{}, fmt.Errorf("%w: expiry_time must not be negative", ErrInvalidRegistrationToken)
	}

	stored := storage.RegistrationToken{
		Token:       token,
		UsesAllowed: req.UsesAllowed,
		ExpiresAtMS: req.ExpiresAtMS,
		CreatedAtMS: t.nowMS(),
	}
	if err := t.store.CreateRegistrationToken(ctx, stored); err != nil {
		return storage.RegistrationToken{}, err
	}
	return stored, nil
}

// RegistrationToken reads one token made through the admin API.
func (t *RegistrationTokens) RegistrationToken(ctx context.Context, token string) (storage.RegistrationToken, error) {
	if t.store == nil {
		return storage.RegistrationToken{}, ErrNoSuchRegistrationToken
	}
	return t.store.RegistrationToken(ctx, token)
}

// List lists the tokens made through the admin API. The config file's are
// not listed: they are in the config file.
func (t *RegistrationTokens) List(ctx context.Context) ([]storage.RegistrationToken, error) {
	if t.store == nil {
		return []storage.RegistrationToken{}, nil
	}
	return t.store.RegistrationTokens(ctx)
}

// Update replaces a token's use limit and expiry.
func (t *RegistrationTokens) Update(ctx context.Context, token string, usesAllowed *int, expiresAtMS int64) (storage.RegistrationToken, error) {
	if t.store == nil {
		return storage.RegistrationToken{}, ErrNoSuchRegistrationToken
	}
	if usesAllowed != nil && *usesAllowed < 0 {
		return storage.RegistrationToken{}, fmt.Errorf("%w: uses_allowed must not be negative", ErrInvalidRegistrationToken)
	}
	if expiresAtMS < 0 {
		return storage.RegistrationToken{}, fmt.Errorf("%w: expiry_time must not be negative", ErrInvalidRegistrationToken)
	}
	if err := t.store.UpdateRegistrationToken(ctx, token, usesAllowed, expiresAtMS); err != nil {
		return storage.RegistrationToken{}, err
	}
	return t.store.RegistrationToken(ctx, token)
}

// Delete removes a token. Accounts registered with it are not affected.
func (t *RegistrationTokens) Delete(ctx context.Context, token string) error {
	if t.store == nil {
		return ErrNoSuchRegistrationToken
	}
	return t.store.DeleteRegistrationToken(ctx, token)
}

// wellFormedToken applies the spec's grammar: 1 to 64 characters from
// [A-Za-z0-9._~-].
func wellFormedToken(token string) bool {
	if token == "" || len(token) > MaxRegistrationTokenLength {
		return false
	}
	for i := range len(token) {
		switch c := token[i]; {
		case 'A' <= c && c <= 'Z', 'a' <= c && c <= 'z', '0' <= c && c <= '9':
		case c == '.', c == '_', c == '~', c == '-':
		default:
			return false
		}
	}
	return true
}

// generateRegistrationToken makes a random token. Bytes that would bias
// the modulo are thrown away, as for media IDs.
func generateRegistrationToken(length int) (string, error) {
	const accept = 256 - 256%len(registrationTokenAlphabet)
	out := make([]byte, 0, length)
	buf := make([]byte, length)
	for len(out) < length {
		if _, err := rand.Read(buf); err != nil {
			return "", fmt.Errorf("generate a registration token: %w", err)
		}
		for _, c := range buf {
			if int(c) < accept && len(out) < length {
				out = append(out, registrationTokenAlphabet[int(c)%len(registrationTokenAlphabet)])
			}
		}
	}
	return string(out), nil
}