// Copyright 2021 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

// Package sdk contains a high-level API for embedding OPA inside of Go programs.
package sdk

import (
	"bytes"
	"cmp"
	"context"
	"crypto/rand"
	"fmt"
	"maps"
	"sync"
	"time"

	"github.com/open-policy-agent/opa/internal/ref"
	"github.com/open-policy-agent/opa/internal/uuid"
	"github.com/open-policy-agent/opa/v1/ast"
	"github.com/open-policy-agent/opa/v1/bundle"
	"github.com/open-policy-agent/opa/v1/hooks"
	"github.com/open-policy-agent/opa/v1/logging"
	"github.com/open-policy-agent/opa/v1/metrics"
	"github.com/open-policy-agent/opa/v1/plugins"
	"github.com/open-policy-agent/opa/v1/plugins/discovery"
	"github.com/open-policy-agent/opa/v1/plugins/logs"
	"github.com/open-policy-agent/opa/v1/rego"
	"github.com/open-policy-agent/opa/v1/runtime/info"
	"github.com/open-policy-agent/opa/v1/server"
	"github.com/open-policy-agent/opa/v1/server/types"
	"github.com/open-policy-agent/opa/v1/storage"
	"github.com/open-policy-agent/opa/v1/topdown"
	"github.com/open-policy-agent/opa/v1/topdown/builtins"
	"github.com/open-policy-agent/opa/v1/topdown/cache"
	"github.com/open-policy-agent/opa/v1/topdown/print"
	"github.com/open-policy-agent/opa/v1/util"
	"github.com/open-policy-agent/opa/v1/version"
)

// OPA represents an instance of the policy engine. OPA can be started with
// several options that control configuration, logging, and lifecycle.
type OPA struct {
	id          string
	state       *state
	mtx         sync.Mutex
	logger      logging.Logger
	console     logging.Logger
	plugins     map[string]plugins.Factory
	store       storage.Store
	hooks       hooks.Hooks
	config      []byte
	regoVersion ast.RegoVersion
	managerOpts []func(*plugins.Manager)
}

type state struct {
	manager                     *plugins.Manager
	interQueryBuiltinCache      cache.InterQueryCache
	interQueryBuiltinValueCache cache.InterQueryValueCache
	queryCache                  *queryCache
}

// New returns a new OPA object. This function should minimally be called with
// options that specify an OPA configuration file.
func New(ctx context.Context, nopts Options) (*OPA, error) {
	defaultOptsMtx.Lock()
	opts := Options{
		Plugins: make(map[string]plugins.Factory, len(nopts.Plugins)+len(defaultOptions.Plugins)),
	}
	opts.ID = cmp.Or(nopts.ID, defaultOptions.ID)
	opts.Config = cmp.Or(nopts.Config, defaultOptions.Config)
	opts.ConsoleLogger = cmp.Or(nopts.ConsoleLogger, defaultOptions.ConsoleLogger)

	opts.V1Compatible = nopts.V1Compatible
	opts.V0Compatible = nopts.V0Compatible
	opts.RegoVersion = cmp.Or(nopts.RegoVersion, defaultOptions.RegoVersion)

	opts.ManagerOpts = append(opts.ManagerOpts, defaultOptions.ManagerOpts...)
	opts.ManagerOpts = append(opts.ManagerOpts, nopts.ManagerOpts...)

	opts.Logger = cmp.Or(nopts.Logger, defaultOptions.Logger)
	hs := make([]hooks.Hook, 0, opts.Hooks.Len()+nopts.Hooks.Len())
	opts.Hooks.Each(func(h hooks.Hook) {
		hs = append(hs, h)
	})
	nopts.Hooks.Each(func(h hooks.Hook) {
		hs = append(hs, h)
	})
	opts.Hooks = hooks.New(hs...)

	maps.Copy(opts.Plugins, defaultOptions.Plugins)
	maps.Copy(opts.Plugins, nopts.Plugins)

	opts.Store = cmp.Or(nopts.Store, defaultOptions.Store)
	opts.Ready = cmp.Or(nopts.Ready, defaultOptions.Ready)
	defaultOptsMtx.Unlock()

	var err error

	id := opts.ID
	if id == "" {
		id, err = uuid.New(rand.Reader)
		if err != nil {
			return nil, err
		}
	}

	if err := opts.init(); err != nil {
		return nil, err
	}

	opa := &OPA{
		id:    id,
		store: opts.Store,
		hooks: opts.Hooks,
		state: &state{
			queryCache: newQueryCache(),
		},
	}

	opa.config = opts.config
	opa.logger = opts.Logger
	opa.console = opts.ConsoleLogger
	opa.plugins = opts.Plugins
	opa.managerOpts = opts.ManagerOpts

	opa.regoVersion = opts.regoVersion()

	return opa, opa.configure(ctx, opa.config, opts.Ready, opts.block)
}

// Plugin returns the named plugin. If the plugin does not exist, this function
// returns nil.
func (opa *OPA) Plugin(name string) plugins.Plugin {
	// Release opa.mtx before calling into the manager to avoid inverting the
	// lock order used by the manager's onCommit callback (see #8873).
	opa.mtx.Lock()
	mgr := opa.state.manager
	opa.mtx.Unlock()

	return mgr.Plugin(name)
}

// Configure updates the configuration of the OPA in-place. This function should
// be called in response to configuration updates in the environment. This
// function is atomic. If the configuration update cannot be successfully
// applied, the old configuration will remain intact.
func (opa *OPA) Configure(ctx context.Context, opts ConfigOptions) error {
	if err := opts.init(); err != nil {
		return err
	}

	// NOTE(tsandall): In future we could be more intelligent about
	// re-configuration and avoid expensive background processing.
	opa.mtx.Lock()
	equal := bytes.Equal(opts.config, opa.config)
	opa.mtx.Unlock()

	if equal {
		close(opts.Ready)
		return nil
	}

	return opa.configure(ctx, opts.config, opts.Ready, opts.block)
}

func (opa *OPA) configure(ctx context.Context, bs []byte, ready chan struct{}, block bool) error {
	runtimeInfo, err := info.NewWithOptions(info.Options{Config: opa.config})
	if err != nil {
		return err
	}

	opts := append([]func(*plugins.Manager){
		plugins.Info(runtimeInfo),
		plugins.Logger(opa.logger),
		plugins.ConsoleLogger(opa.console),
		plugins.WithParserOptions(ast.ParserOptions{ProcessAnnotation: true, RegoVersion: opa.regoVersion}),
		plugins.EnablePrintStatements(opa.logger.GetLevel() >= logging.Info),
		plugins.PrintHook(loggingPrintHook{logger: opa.logger}),
		plugins.WithHooks(opa.hooks),
	}, opa.managerOpts...)

	// Plumb in storage for external bundle activation plugin, if registered with bundle.RegisterStore,
	// unless the user has passed their own store already.
	var store storage.Store
	switch {
	case opa.store != nil:
		store = opa.store
	case bundle.BundleExtStore != nil:
		store = bundle.BundleExtStore()
	}
	manager, err := plugins.New(
		bs,
		opa.id,
		store,
		opts...,
	)
	if err != nil {
		return err
	}

	manager.RegisterCompilerTrigger(func(storage.Transaction) {
		opa.mtx.Lock()
		opa.state.queryCache.Clear()
		opa.mtx.Unlock()
	})

	manager.RegisterPluginStatusListener("sdk", func(status map[string]*plugins.Status) {
		select {
		case <-ready:
			return
		default:
		}
		// NOTE(tsandall): we do not include a special case for the discovery
		// plugin. If the discovery plugin is the only plugin and it goes ready,
		// then OPA will be considered ready. The discovery plugin only goes ready
		// _after_ it has successfully processed a discovery bundle. During
		// discovery bundle processing, other plugins will register so their states
		// will be accounted for. If a discovery bundle did not enable any other
		// plugins (bundles, etc.) the OPA will still be operational.
		for _, s := range status {
			if s.State != plugins.StateOK {
				return
			}
		}

		close(ready)
	})

	var bootConfig map[string]any
	err = util.Unmarshal(opa.config, &bootConfig)
	if err != nil {
		return err
	}

	d, err := discovery.New(manager,
		discovery.Factories(opa.plugins),
		discovery.Hooks(opa.hooks),
		discovery.BootConfig(bootConfig),
	)
	if err != nil {
		return err
	}

	manager.Register(discovery.Name, d)

	opa.mtx.Lock()

	// NOTE(tsandall): there is no return value from Stop() and it could block
	// on async operations (e.g., decision log uploading) so defer the call to
	// another goroutine.
	//
	// TODO(tsandall): if we need to block on operations like decision log
	// uploading, perhaps we could rely on a manual trigger.
	previousManager := opa.state.manager
	go func() {
		if previousManager != nil {
			previousManager.Stop(ctx)
		}
	}()

	opa.state.manager = manager
	opa.state.queryCache.Clear()
	opa.state.interQueryBuiltinCache = cache.NewInterQueryCacheWithContext(ctx, manager.InterQueryBuiltinCacheConfig())
	opa.state.interQueryBuiltinValueCache = cache.NewInterQueryValueCache(ctx, manager.InterQueryBuiltinCacheConfig())
	opa.config = bs

	opa.mtx.Unlock()

	if err := manager.Start(ctx); err != nil {
		return err
	}

	// Resolve the buffered logger: flush to logger plugin if configured,
	// otherwise discard (no fallback).
	opa.logger = manager.ResolveBufferedLogger(nil)

	if block {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ready:
		}
	}

	return nil
}

// Stop closes the OPA. The OPA cannot be restarted.
func (opa *OPA) Stop(ctx context.Context) {
	opa.mtx.Lock()
	mgr := opa.state.manager
	opa.mtx.Unlock()

	if mgr != nil {
		mgr.Stop(ctx)
	}
}

// Decision returns a named decision. This function is threadsafe.
func (opa *OPA) Decision(ctx context.Context, options DecisionOptions) (*DecisionResult, error) {
	record := server.Info{
		Timestamp:      options.Now,
		Path:           options.Path,
		Input:          &options.Input,
		NDBuiltinCache: &options.NDBCache,
		Metrics:        options.Metrics,
		DecisionID:     options.DecisionID,
	}

	// Only use non-deterministic builtins cache if it's available.
	var ndbc builtins.NDBCache
	if options.NDBCache != nil {
		if v, ok := options.NDBCache.(builtins.NDBCache); ok {
			ndbc = v
		}
	}

	// TODO: make extractor configurable via SDK options
	tracker := &topdown.EvaluatedRuleTracker{}

	result, err := opa.executeTransaction(
		ctx,
		&record,
		func(s state, result *DecisionResult) {
			result.Result, result.Provenance, record.InputAST, record.Bundles, record.Error = evaluate(ctx, evalArgs{
				runtime:                     s.manager.Info,
				printHook:                   s.manager.PrintHook(),
				compiler:                    s.manager.GetCompiler(),
				store:                       s.manager.Store,
				queryCache:                  s.queryCache,
				interQueryCache:             s.interQueryBuiltinCache,
				interQueryBuiltinValueCache: s.interQueryBuiltinValueCache,
				ndbcache:                    ndbc,
				txn:                         record.Txn,
				now:                         record.Timestamp,
				path:                        record.Path,
				input:                       *record.Input,
				m:                           record.Metrics,
				strictBuiltinErrors:         options.StrictBuiltinErrors,
				tracer:                      options.Tracer,
				profiler:                    options.Profiler,
				instrument:                  options.Instrument,
				evaluatedRules:              tracker,
				httpRoundTripper:            options.HTTPRoundTripper,
			})
			if record.Error == nil {
				record.Results = &result.Result
			}
			record.EvaluatedRuleLabels = tracker.Labels
		},
	)
	if err != nil {
		return nil, err
	}

	return result, record.Error
}

// DecisionOptions contains parameters for query evaluation.
type DecisionOptions struct {
	Now                 time.Time           // specifies wallclock time used for time.now_ns(), decision log timestamp, etc.
	Path                string              // specifies name of policy decision to evaluate (e.g., example/allow)
	Input               any                 // specifies value of the input document to evaluate policy with
	NDBCache            any                 // specifies the non-deterministic builtins cache to use for evaluation.
	StrictBuiltinErrors bool                // treat built-in function errors as fatal
	Tracer              topdown.QueryTracer // specifies the tracer to use for evaluation, optional
	Metrics             metrics.Metrics     // specifies the metrics to use for preparing and evaluation, optional
	Profiler            topdown.QueryTracer // specifies the profiler to use, optional
	Instrument          bool                // if true, instrumentation will be enabled
	DecisionID          string              // the identifier for this decision; if not set, a globally unique identifier will be generated

	// HTTPRoundTripper customizes the http.RoundTripper used by http.send
	// during this decision. The provided function receives the http.Transport
	// OPA would otherwise use — which may be nil for plain-HTTP requests
	// without TLS or unix-socket options — and returns the http.RoundTripper
	// to use in its place, typically a wrapper around the received transport.
	// When nil, OPA's default behavior is preserved.
	HTTPRoundTripper topdown.CustomizeRoundTripper
}

// DecisionResult contains the output of query evaluation.
type DecisionResult struct {
	ID         string             // provides the identifier for this decision (which is included in the decision log.)
	Result     any                // provides the output of query evaluation.
	Provenance types.ProvenanceV1 // wraps the bundle build/version information
}

func (opa *OPA) executeTransaction(ctx context.Context, record *server.Info, work func(state, *DecisionResult)) (*DecisionResult, error) {
	record.Metrics = util.Or(record.Metrics, metrics.New)
	record.Metrics.Timer(metrics.SDKDecisionEval).Start()

	if record.DecisionID == "" {
		id, err := uuid.New(rand.Reader)
		if err != nil {
			return nil, err
		}
		record.DecisionID = id
	}

	result := &DecisionResult{ID: record.DecisionID}

	opa.mtx.Lock()
	s := *opa.state
	opa.mtx.Unlock()

	if record.Timestamp.IsZero() {
		record.Timestamp = time.Now().UTC()
	}

	if record.Path == "" {
		record.Path = *s.manager.GetConfig().DefaultDecision
	}

	record.Txn, record.Error = s.manager.Store.NewTransaction(ctx, storage.TransactionParams{})

	if record.Error == nil {
		defer s.manager.Store.Abort(ctx, record.Txn)
		work(s, result)
	}

	record.Metrics.Timer(metrics.SDKDecisionEval).Stop()

	if logger := logs.Lookup(s.manager); logger != nil {
		// Decision log masking requires the event object to be a map[string]any,
		// or a []any, and all internal objects referenced in the mask to be
		// similarly generic. Convert the input AST back into a JSON-representation to
		// ensure decision logging will work if the input Go type does not fit these requirements.
		if record.InputAST != nil {
			asJSON, err := ast.JSON(record.InputAST)
			if err != nil {
				return nil, err
			}
			*record.Input = asJSON
		}
		if err := logger.Log(ctx, record); err != nil {
			return result, fmt.Errorf("decision log: %w", err)
		}
	}
	return result, nil
}

// Partial returns a named decision. This function is threadsafe.
// Note(philipc): The NDBCache is unused here, because non-deterministic
// builtins are not run during partial evaluation.
func (opa *OPA) Partial(ctx context.Context, options PartialOptions) (*PartialResult, error) {
	if options.Mapper == nil {
		options.Mapper = &RawMapper{}
	}

	record := server.Info{
		Timestamp:  options.Now,
		Input:      &options.Input,
		Query:      options.Query,
		Metrics:    options.Metrics,
		DecisionID: options.DecisionID,
	}

	var provenance types.ProvenanceV1

	var pq *rego.PartialQueries
	decision, err := opa.executeTransaction(
		ctx,
		&record,
		func(s state, result *DecisionResult) {
			pq, provenance, record.InputAST, record.Bundles, record.Error = partial(ctx, partialEvalArgs{
				runtime:             s.manager.Info,
				printHook:           s.manager.PrintHook(),
				compiler:            s.manager.GetCompiler(),
				store:               s.manager.Store,
				txn:                 record.Txn,
				now:                 record.Timestamp,
				query:               record.Query,
				unknowns:            options.Unknowns,
				input:               *record.Input,
				m:                   record.Metrics,
				strictBuiltinErrors: options.StrictBuiltinErrors,
				tracer:              options.Tracer,
				profiler:            options.Profiler,
				instrument:          options.Instrument,
			})
			if record.Error == nil {
				result.Result, record.Error = options.Mapper.MapResults(pq)
				var pqAst any
				if record.Error == nil {
					var mappedResults any
					mappedResults, record.Error = options.Mapper.ResultToJSON(result.Result)
					record.MappedResults = &mappedResults
					pqAst = pq
					record.Results = &pqAst
				}
			}
		},
	)
	if err != nil {
		return nil, err
	}

	return &PartialResult{
		ID:         decision.ID,
		Result:     decision.Result,
		AST:        pq,
		Provenance: provenance,
	}, record.Error
}

type PartialQueryMapper interface {
	// The first interface being returned is the type that will be used for further processing
	MapResults(pq *rego.PartialQueries) (any, error)
	// This should be able to take the Result object from MapResults and return a type that can be logged as JSON
	ResultToJSON(result any) (any, error)
}

// PartialOptions contains parameters for partial query evaluation.
type PartialOptions struct {
	Now                 time.Time           // specifies wallclock time used for time.now_ns(), decision log timestamp, etc.
	Input               any                 // specifies value of the input document to evaluate policy with
	Query               string              // specifies the query to be partially evaluated
	Unknowns            []string            // specifies the unknown elements of the policy
	Mapper              PartialQueryMapper  // specifies the mapper to use when processing results
	StrictBuiltinErrors bool                // treat built-in function errors as fatal
	Tracer              topdown.QueryTracer // specifies the tracer to use for evaluation, optional
	Metrics             metrics.Metrics     // specifies the metrics to use for preparing and evaluation, optional
	Profiler            topdown.QueryTracer // specifies the profiler to use, optional
	Instrument          bool                // if true, instrumentation will be enabled
	DecisionID          string              // the identifier for this decision; if not set, a globally unique identifier will be generated
}

type PartialResult struct {
	ID         string               // decision ID
	Result     any                  // mapped result
	AST        *rego.PartialQueries // raw result
	Provenance types.ProvenanceV1   // wraps the bundle build/version information
}

// Error represents an internal error in the SDK.
type Error struct {
	Code    string `json:"code"`
	Message string `json:"message,omitempty"`
}

func (err *Error) Error() string {
	return fmt.Sprintf("%v: %v", err.Code, err.Message)
}

const (
	// UndefinedErr indicates that the queried decision was undefined.
	UndefinedErr = "opa_undefined_error"
)

func undefinedDecisionErr(path string) *Error {
	return &Error{
		Code:    UndefinedErr,
		Message: path + " decision was undefined",
	}
}

// IsUndefinedErr returns true of the err represents an undefined decision error.
func IsUndefinedErr(err error) bool {
	actual, ok := err.(*Error)
	return ok && actual.Code == UndefinedErr
}

type evalArgs struct {
	runtime                     *ast.Term
	printHook                   print.Hook
	compiler                    *ast.Compiler
	store                       storage.Store
	txn                         storage.Transaction
	queryCache                  *queryCache
	interQueryCache             cache.InterQueryCache
	interQueryBuiltinValueCache cache.InterQueryValueCache
	now                         time.Time
	path                        string
	input                       any
	ndbcache                    builtins.NDBCache
	m                           metrics.Metrics
	strictBuiltinErrors         bool
	tracer                      topdown.QueryTracer
	profiler                    topdown.QueryTracer
	instrument                  bool
	evaluatedRules              *topdown.EvaluatedRuleTracker
	httpRoundTripper            topdown.CustomizeRoundTripper
}

func evaluate(ctx context.Context, args evalArgs) (any, types.ProvenanceV1, ast.Value, map[string]server.BundleInfo, error) {
	provenance := types.ProvenanceV1{
		Version:   version.Version,
		Vcs:       version.Vcs,
		Timestamp: version.Timestamp,
		Hostname:  version.Hostname,
		Bundles:   make(map[string]types.ProvenanceBundleV1),
	}
	bundles, err := bundles(ctx, args.store, args.txn)
	if err != nil {
		return nil, provenance, nil, nil, err
	}
	for b, info := range bundles {
		provenance.Bundles[b] = types.ProvenanceBundleV1{
			Revision: info.Revision,
		}
	}

	r, err := ref.ParseDataPath(args.path)
	if err != nil {
		return nil, provenance, nil, bundles, err
	}

	pq, err := args.queryCache.Get(r.String(), func(query string) (*rego.PreparedEvalQuery, error) {
		opts := []func(*rego.Rego){
			rego.Time(args.now),
			rego.Metrics(args.m),
			rego.Query(query),
			rego.Compiler(args.compiler),
			rego.Store(args.store),
			rego.Transaction(args.txn),
			rego.PrintHook(args.printHook),
			rego.StrictBuiltinErrors(args.strictBuiltinErrors),
			rego.Instrument(args.instrument),
			rego.Runtime(args.runtime),
		}
		if args.evaluatedRules != nil {
			opts = append(opts, rego.EvaluatedRuleTracker(args.evaluatedRules))
		}
		pq, err := rego.New(opts...).PrepareForEval(ctx)
		if err != nil {
			return nil, err
		}
		return &pq, err
	})
	if err != nil {
		return nil, provenance, nil, bundles, err
	}

	inputAST, err := ast.InterfaceToValue(args.input)
	if err != nil {
		return nil, provenance, nil, bundles, err
	}

	rs, err := pq.Eval(
		ctx,
		rego.EvalTime(args.now),
		rego.EvalParsedInput(inputAST),
		rego.EvalTransaction(args.txn),
		rego.EvalInterQueryBuiltinCache(args.interQueryCache),
		rego.EvalInterQueryBuiltinValueCache(args.interQueryBuiltinValueCache),
		rego.EvalNDBuiltinCache(args.ndbcache),
		rego.EvalQueryTracer(args.tracer),
		rego.EvalMetrics(args.m),
		rego.EvalQueryTracer(args.profiler),
		rego.EvalInstrument(args.instrument),
		rego.EvalHTTPRoundTripper(args.httpRoundTripper),
	)
	if err != nil {
		return nil, provenance, inputAST, bundles, err
	} else if len(rs) == 0 {
		return nil, provenance, inputAST, bundles, undefinedDecisionErr(args.path)
	}

	return rs[0].Expressions[0].Value, provenance, inputAST, bundles, nil
}

type partialEvalArgs struct {
	runtime             *ast.Term
	compiler            *ast.Compiler
	printHook           print.Hook
	store               storage.Store
	txn                 storage.Transaction
	unknowns            []string
	query               string
	now                 time.Time
	input               any
	m                   metrics.Metrics
	strictBuiltinErrors bool
	tracer              topdown.QueryTracer
	profiler            topdown.QueryTracer
	instrument          bool
}

func partial(ctx context.Context, args partialEvalArgs) (*rego.PartialQueries, types.ProvenanceV1, ast.Value, map[string]server.BundleInfo, error) {
	provenance := types.ProvenanceV1{
		Version: version.Version,
		Bundles: make(map[string]types.ProvenanceBundleV1),
	}

	bundles, err := bundles(ctx, args.store, args.txn)
	if err != nil {
		return nil, provenance, nil, nil, err
	}
	for b, info := range bundles {
		provenance.Bundles[b] = types.ProvenanceBundleV1{
			Revision: info.Revision,
		}
	}

	inputAST, err := ast.InterfaceToValue(args.input)
	if err != nil {
		return nil, provenance, nil, bundles, err
	}
	re := rego.New(
		rego.Time(args.now),
		rego.Metrics(args.m),
		rego.Store(args.store),
		rego.Compiler(args.compiler),
		rego.Transaction(args.txn),
		rego.Runtime(args.runtime),
		rego.Input(args.input),
		rego.Query(args.query),
		rego.Unknowns(args.unknowns),
		rego.PrintHook(args.printHook),
		rego.StrictBuiltinErrors(args.strictBuiltinErrors),
		rego.QueryTracer(args.tracer),
		rego.QueryTracer(args.profiler),
		rego.Instrument(args.instrument),
	)

	pq, err := re.Partial(ctx)
	if err != nil {
		return nil, provenance, nil, bundles, err
	}
	return pq, provenance, inputAST, bundles, err
}

type queryCache struct {
	sync.Mutex
	cache map[string]*rego.PreparedEvalQuery
}

func newQueryCache() *queryCache {
	return &queryCache{cache: map[string]*rego.PreparedEvalQuery{}}
}

func (qc *queryCache) Get(key string, orElse func(string) (*rego.PreparedEvalQuery, error)) (*rego.PreparedEvalQuery, error) {
	qc.Lock()
	defer qc.Unlock()

	result, ok := qc.cache[key]
	if ok {
		return result, nil
	}

	result, err := orElse(key)
	if err != nil {
		return nil, err
	}

	qc.cache[key] = result
	return result, nil
}

func (qc *queryCache) Clear() {
	qc.Lock()
	defer qc.Unlock()

	qc.cache = make(map[string]*rego.PreparedEvalQuery)
}

func bundles(ctx context.Context, store storage.Store, txn storage.Transaction) (map[string]server.BundleInfo, error) {
	bundles := map[string]server.BundleInfo{}
	names, err := bundle.ReadBundleNamesFromStore(ctx, store, txn)
	if err != nil && !storage.IsNotFound(err) {
		return nil, fmt.Errorf("failed to read bundle names: %w", err)
	}
	for _, name := range names {
		r, err := bundle.ReadBundleRevisionFromStore(ctx, store, txn, name)
		if err != nil {
			return nil, fmt.Errorf("failed to read bundle revisions: %w", err)
		}
		bundles[name] = server.BundleInfo{Revision: r}
	}
	return bundles, nil
}

type loggingPrintHook struct {
	logger logging.Logger
}

func (h loggingPrintHook) Print(pctx print.Context, msg string) error {
	h.logger.WithFields(map[string]any{"line": pctx.Location.String()}).Info("%s", msg)
	return nil
}