package tokens

import (
	"context"
	"crypto/subtle"
	"fmt"
	"strings"

	"atlas9.dev/c/core"
)

// VerifiedData wraps user data with a verification code.
// Used as the data type for the inner store of a SecureStore.
type VerifiedData[T any] struct {
	Code string
	Data T
}

// SecureStore wraps a Store to protect against timing attacks.
// It splits the token into a key (for database lookup) and a verification code
// (compared using crypto/subtle.ConstantTimeCompare).
//
// The combined token given to the user is "key.code". On Get,
// the key is used for the DB lookup and the code is compared in constant time.
//
// SecureStore implements Store[T] and can be used as a drop-in replacement.
type SecureStore[T any] struct {
	store   Store[VerifiedData[T]]
	codeGen KeyGen
}

// NewSecureStore creates a SecureStore that wraps an inner store.
// The codeGen function generates the verification code portion of the token.
func NewSecureStore[T any](store Store[VerifiedData[T]], codeGen KeyGen) *SecureStore[T] {
	return &SecureStore[T]{store: store, codeGen: codeGen}
}

// Put stores data under the given key with a generated verification code.
// The returned Token.Key is the combined "key.code" string to give to the user.
func (s *SecureStore[T]) Put(ctx context.Context, key string, data T) (*Token[T], error) {
	code, err := s.codeGen()
	if err != nil {
		return nil, fmt.Errorf("generating verification code: %w", err)
	}

	tok, err := s.store.Put(ctx, key, VerifiedData[T]{
		Code: code,
		Data: data,
	})
	if err != nil {
		return nil, err
	}

	return &Token[T]{
		Key:       tok.Key + "." + code,
		ExpiresAt: tok.ExpiresAt,
		Data:      data,
	}, nil
}

// Get splits the combined token, looks up by key, and verifies the code
// using constant-time comparison. Returns core.ErrNotFound if the token is
// invalid, expired, or the verification code doesn't match.
func (s *SecureStore[T]) Get(ctx context.Context, combined string) (*Token[T], error) {
	key, code, ok := strings.Cut(combined, ".")
	if !ok {
		return nil, core.ErrNotFound
	}

	tok, err := s.store.Get(ctx, key)
	if err != nil {
		return nil, err
	}

	if subtle.ConstantTimeCompare([]byte(tok.Data.Code), []byte(code)) != 1 {
		return nil, core.ErrNotFound
	}

	return &Token[T]{
		Key:       combined,
		ExpiresAt: tok.ExpiresAt,
		Data:      tok.Data.Data,
	}, nil
}

// Delete splits the combined token and deletes by the raw key.
func (s *SecureStore[T]) Delete(ctx context.Context, combined string) error {
	key, _, ok := strings.Cut(combined, ".")
	if !ok {
		return core.ErrNotFound
	}
	return s.store.Delete(ctx, key)
}