package executor

import "context"

// KeyspaceSetter updates the current keyspace for a session.
type KeyspaceSetter interface {
	SetKeyspace(keyspace string)
}

// KeyspaceGetter reads the current keyspace for a session.
type KeyspaceGetter interface {
	GetKeyspace() string
}

// KeyspaceState provides access to the session keyspace.
type KeyspaceState interface {
	KeyspaceSetter
	KeyspaceGetter
}

type keyspaceSetterKey struct{}
type allowMissingKeyspaceKey struct{}

// WithKeyspaceState stores keyspace state in the context for executor use.
func WithKeyspaceState(ctx context.Context, state KeyspaceState) context.Context {
	if state == nil {
		return ctx
	}
	return context.WithValue(ctx, keyspaceSetterKey{}, state)
}

// WithKeyspaceSetter stores a keyspace setter in the context for executor use.
func WithKeyspaceSetter(ctx context.Context, setter KeyspaceSetter) context.Context {
	if state, ok := setter.(KeyspaceState); ok {
		return WithKeyspaceState(ctx, state)
	}
	if setter == nil {
		return ctx
	}
	return context.WithValue(ctx, keyspaceSetterKey{}, setter)
}

// WithAllowMissingKeyspace allows PREPARE to skip keyspace validation.
func WithAllowMissingKeyspace(ctx context.Context) context.Context {
	return context.WithValue(ctx, allowMissingKeyspaceKey{}, true)
}

func keyspaceStateFromContext(ctx context.Context) KeyspaceState {
	if v := ctx.Value(keyspaceSetterKey{}); v != nil {
		if state, ok := v.(KeyspaceState); ok {
			return state
		}
	}
	return nil
}

func allowMissingKeyspaceFromContext(ctx context.Context) bool {
	if v := ctx.Value(allowMissingKeyspaceKey{}); v != nil {
		if allowed, ok := v.(bool); ok {
			return allowed
		}
	}
	return false
}