package integration

import (
	"context"

	"forge.cadoles.com/cadoles/bouncer/internal/jwk"
	"github.com/pkg/errors"
)

var (
	ErrNotFound = errors.New("not found")
)

type contextKey string

const (
	ctxPublicKeySet contextKey = "public-key-set"
	ctxPrivateKey   contextKey = "private-key"
)

func CtxPublicKeySet(ctx context.Context) (jwk.Set, error) {
	return ctxValue[jwk.Set](ctx, ctxPublicKeySet)
}

func WithPublicKeySet(ctx context.Context, set jwk.Set) context.Context {
	return context.WithValue(ctx, ctxPublicKeySet, set)
}

func CtxPrivateKey(ctx context.Context) (jwk.Key, error) {
	return ctxValue[jwk.Key](ctx, ctxPrivateKey)
}

func WithPrivateKey(ctx context.Context, key jwk.Key) context.Context {
	return context.WithValue(ctx, ctxPrivateKey, key)
}

func ctxValue[T any](ctx context.Context, key contextKey) (T, error) {
	raw := ctx.Value(key)
	if raw == nil {
		return *new(T), errors.WithStack(ErrNotFound)
	}

	value, ok := raw.(T)
	if !ok {
		return *new(T), errors.Errorf("unexpected value type '%T'", raw)
	}

	return value, nil
}