package topdown

import (
	"cmp"
	"context"
	"crypto/rand"
	"io"
	"slices"
	"time"

	"github.com/open-policy-agent/opa/v1/ast"
	"github.com/open-policy-agent/opa/v1/metrics"
	"github.com/open-policy-agent/opa/v1/resolver"
	"github.com/open-policy-agent/opa/v1/storage"
	"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/copypropagation"
	"github.com/open-policy-agent/opa/v1/topdown/print"
	"github.com/open-policy-agent/opa/v1/tracing"
	"github.com/open-policy-agent/opa/v1/util"
)

// QueryResultSet represents a collection of results returned by a query.
type QueryResultSet []QueryResult

// QueryResult represents a single result returned by a query. The result
// contains bindings for all variables that appear in the query.
type QueryResult map[ast.Var]*ast.Term

// Query provides a configurable interface for performing query evaluation.
type Query struct {
	seed                        io.Reader
	time                        time.Time
	cancel                      Cancel
	query                       ast.Body
	queryCompiler               ast.QueryCompiler
	compiler                    *ast.Compiler
	store                       storage.Store
	txn                         storage.Transaction
	input                       *ast.Term
	external                    *resolverTrie
	tracers                     []QueryTracer
	plugTraceVars               bool
	unknowns                    []*ast.Term
	partialNamespace            string
	skipSaveNamespace           bool
	metrics                     metrics.Metrics
	instr                       *Instrumentation
	disableInlining             []ast.Ref
	shallowInlining             bool
	nondeterministicBuiltins    bool
	genvarprefix                string
	runtime                     *ast.Term
	builtins                    map[string]*Builtin
	indexing                    bool
	earlyExit                   bool
	interQueryBuiltinCache      cache.InterQueryCache
	interQueryBuiltinValueCache cache.InterQueryValueCache
	ndBuiltinCache              builtins.NDBCache
	strictBuiltinErrors         bool
	stackTraces                 bool
	builtinErrorList            *[]Error
	strictObjects               bool
	roundTripper                CustomizeRoundTripper
	printHook                   print.Hook
	tracingOpts                 tracing.Options
	virtualCache                VirtualCache
	baseCache                   BaseCache
	requestMetadata             map[string]any
	responseMetadata            map[string]any
	evaluated                   *EvaluatedRuleTracker
}

// Builtin represents a built-in function that queries can call.
type Builtin struct {
	Decl *ast.Builtin
	Func BuiltinFunc
}

// NewQuery returns a new Query object that can be run.
func NewQuery(query ast.Body) *Query {
	return &Query{
		query:        query,
		genvarprefix: ast.WildcardPrefix,
		indexing:     true,
		earlyExit:    true,
	}
}

// WithQueryCompiler sets the queryCompiler used for the query.
func (q *Query) WithQueryCompiler(queryCompiler ast.QueryCompiler) *Query {
	q.queryCompiler = queryCompiler
	return q
}

// WithCompiler sets the compiler to use for the query.
func (q *Query) WithCompiler(compiler *ast.Compiler) *Query {
	q.compiler = compiler
	return q
}

// WithStore sets the store to use for the query.
func (q *Query) WithStore(store storage.Store) *Query {
	q.store = store
	return q
}

// WithTransaction sets the transaction to use for the query. All queries
// should be performed over a consistent snapshot of the storage layer.
func (q *Query) WithTransaction(txn storage.Transaction) *Query {
	q.txn = txn
	return q
}

// WithCancel sets the cancellation object to use for the query. Set this if
// you need to abort queries based on a deadline. This is optional.
func (q *Query) WithCancel(cancel Cancel) *Query {
	q.cancel = cancel
	return q
}

// WithInput sets the input object to use for the query. References rooted at
// input will be evaluated against this value. This is optional.
func (q *Query) WithInput(input *ast.Term) *Query {
	q.input = input
	return q
}

// WithTracer adds a query tracer to use during evaluation. This is optional.
//
// Deprecated: Use WithQueryTracer instead.
func (q *Query) WithTracer(tracer Tracer) *Query {
	qt, ok := tracer.(QueryTracer)
	if !ok {
		qt = WrapLegacyTracer(tracer)
	}
	return q.WithQueryTracer(qt)
}

// WithQueryTracer adds a query tracer to use during evaluation. This is optional.
// Disabled QueryTracers will be ignored.
func (q *Query) WithQueryTracer(tracer QueryTracer) *Query {
	if tracer == nil || !tracer.Enabled() {
		return q
	}

	q.tracers = append(q.tracers, tracer)

	// If *any* of the tracers require local variable metadata we need to
	// enabled plugging local trace variables.
	conf := tracer.Config()
	if conf.PlugLocalVars {
		q.plugTraceVars = true
	}

	return q
}

// WithMetrics sets the metrics collection to add evaluation metrics to. This
// is optional.
func (q *Query) WithMetrics(m metrics.Metrics) *Query {
	q.metrics = m
	return q
}

// WithInstrumentation sets the instrumentation configuration to enable on the
// evaluation process. By default, instrumentation is turned off.
func (q *Query) WithInstrumentation(instr *Instrumentation) *Query {
	q.instr = instr
	return q
}

// WithUnknowns sets the initial set of variables or references to treat as
// unknown during query evaluation. This is required for partial evaluation.
func (q *Query) WithUnknowns(terms []*ast.Term) *Query {
	q.unknowns = terms
	return q
}

// WithPartialNamespace sets the namespace to use for supporting rules
// generated as part of the partial evaluation process. The ns value must be a
// valid package path component.
func (q *Query) WithPartialNamespace(ns string) *Query {
	q.partialNamespace = ns
	return q
}

// WithSkipPartialNamespace disables namespacing of saved support rules that are generated
// from the original policy (rules which are completely synthetic are still namespaced.)
func (q *Query) WithSkipPartialNamespace(yes bool) *Query {
	q.skipSaveNamespace = yes
	return q
}

// WithDisableInlining adds a set of paths to the query that should be excluded from
// inlining. Inlining during partial evaluation can be expensive in some cases
// (e.g., when a cross-product is computed.) Disabling inlining avoids expensive
// computation at the cost of generating support rules.
func (q *Query) WithDisableInlining(paths []ast.Ref) *Query {
	q.disableInlining = paths
	return q
}

// WithShallowInlining disables aggressive inlining performed during partial evaluation.
// When shallow inlining is enabled rules that depend (transitively) on unknowns are not inlined.
// Only rules/values that are completely known will be inlined.
func (q *Query) WithShallowInlining(yes bool) *Query {
	q.shallowInlining = yes
	return q
}

// WithRuntime sets the runtime data to execute the query with. The runtime data
// can be returned by the `opa.runtime` built-in function.
func (q *Query) WithRuntime(runtime *ast.Term) *Query {
	q.runtime = runtime
	return q
}

// WithBuiltins adds a set of built-in functions that can be called by the
// query.
func (q *Query) WithBuiltins(builtins map[string]*Builtin) *Query {
	q.builtins = builtins
	return q
}

// WithIndexing will enable or disable using rule indexing for the evaluation
// of the query. The default is enabled.
func (q *Query) WithIndexing(enabled bool) *Query {
	q.indexing = enabled
	return q
}

// WithEarlyExit will enable or disable using 'early exit' for the evaluation
// of the query. The default is enabled.
func (q *Query) WithEarlyExit(enabled bool) *Query {
	q.earlyExit = enabled
	return q
}

// WithSeed sets a reader that will seed randomization required by built-in functions.
// If a seed is not provided crypto/rand.Reader is used.
func (q *Query) WithSeed(r io.Reader) *Query {
	q.seed = r
	return q
}

// WithTime sets the time that will be returned by the time.now_ns() built-in function.
func (q *Query) WithTime(x time.Time) *Query {
	q.time = x
	return q
}

// WithInterQueryBuiltinCache sets the inter-query cache that built-in functions can utilize.
func (q *Query) WithInterQueryBuiltinCache(c cache.InterQueryCache) *Query {
	q.interQueryBuiltinCache = c
	return q
}

// WithInterQueryBuiltinValueCache sets the inter-query value cache that built-in functions can utilize.
func (q *Query) WithInterQueryBuiltinValueCache(c cache.InterQueryValueCache) *Query {
	q.interQueryBuiltinValueCache = c
	return q
}

// WithNDBuiltinCache sets the non-deterministic builtin cache.
func (q *Query) WithNDBuiltinCache(c builtins.NDBCache) *Query {
	q.ndBuiltinCache = c
	return q
}

// WithStrictBuiltinErrors tells the evaluator to treat all built-in function errors as fatal errors.
func (q *Query) WithStrictBuiltinErrors(yes bool) *Query {
	q.strictBuiltinErrors = yes
	return q
}

// WithStackTraces tells the evaluator to record the stack of queries being
// evaluated when an error occurred on the returned *Error. The stack is exposed
// as Error.StackTrace and left out of the error message, so callers render it
// themselves.
//
// Off by default because capture is not free: each *Error costs a walk of the
// parent chain and a frame per query on it, resolved against the bindings in
// scope, which on a query collecting one built-in error per row runs from +60%
// to +203% in time (see BenchmarkStackTraceCollectedBuiltinErrors). OPA's own
// CLI and server accept that cost and turn it on.
func (q *Query) WithStackTraces(yes bool) *Query {
	q.stackTraces = yes
	return q
}

// WithBuiltinErrorList supplies a pointer to an Error slice to store built-in function errors
// encountered during evaluation. This error slice can be inspected after evaluation to determine
// which built-in function errors occurred.
func (q *Query) WithBuiltinErrorList(list *[]Error) *Query {
	q.builtinErrorList = list
	return q
}

// WithResolver configures an external resolver to use for the given ref.
func (q *Query) WithResolver(ref ast.Ref, r resolver.Resolver) *Query {
	q.external = util.Or(q.external, newResolverTrie)
	q.external.Put(ref, r)
	return q
}

// WithHTTPRoundTripper configures a custom HTTP transport for built-in functions that make HTTP requests.
func (q *Query) WithHTTPRoundTripper(t CustomizeRoundTripper) *Query {
	q.roundTripper = t
	return q
}

func (q *Query) WithPrintHook(h print.Hook) *Query {
	q.printHook = h
	return q
}

// WithDistributedTracingOpts sets the options to be used by distributed tracing.
func (q *Query) WithDistributedTracingOpts(tr tracing.Options) *Query {
	q.tracingOpts = tr
	return q
}

// WithStrictObjects tells the evaluator to avoid the "lazy object" optimization
// applied when reading objects from the store. It will result in higher memory
// usage and should only be used temporarily while adjusting code that breaks
// because of the optimization.
func (q *Query) WithStrictObjects(yes bool) *Query {
	q.strictObjects = yes
	return q
}

// WithVirtualCache sets the VirtualCache to use during evaluation. This is
// optional, and if not set, the default cache is used.
func (q *Query) WithVirtualCache(vc VirtualCache) *Query {
	q.virtualCache = vc
	return q
}

// WithBaseCache sets the BaseCache to use during evaluation. This is
// optional, and if not set, the default cache is used.
func (q *Query) WithBaseCache(bc BaseCache) *Query {
	q.baseCache = bc
	return q
}

// WithNondeterministicBuiltins causes non-deterministic builtins to be evalued
// during partial evaluation. This is needed to pull in external data, or validate
// a JWT, during PE, so that the result informs what queries are returned.
func (q *Query) WithNondeterministicBuiltins(yes bool) *Query {
	q.nondeterministicBuiltins = yes
	return q
}

// WithRequestMetadata sets arbitrary metadata from the caller that can be
// used by wrapping projects. The data is stored but not directly used by
// OPA's evaluation engine.
func (q *Query) WithRequestMetadata(m map[string]any) *Query {
	q.requestMetadata = m
	return q
}

// WithResponseMetadata sets a map that wrapping projects can populate during
// evaluation to include additional fields in the API response.
func (q *Query) WithResponseMetadata(m map[string]any) *Query {
	q.responseMetadata = m
	return q
}

// WithEvaluatedRuleTracker sets a tracker to record rule identifiers that were
// successfully evaluated.
func (q *Query) WithEvaluatedRuleTracker(t *EvaluatedRuleTracker) *Query {
	q.evaluated = t
	return q
}

// PartialRun executes partial evaluation on the query with respect to unknown
// values. Partial evaluation attempts to evaluate as much of the query as
// possible without requiring values for the unknowns set on the query. The
// result of partial evaluation is a new set of queries that can be evaluated
// once the unknown value is known. In addition to new queries, partial
// evaluation may produce additional support modules that should be used in
// conjunction with the partially evaluated queries.
func (q *Query) PartialRun(ctx context.Context) (partials []ast.Body, support []*ast.Module, err error) {
	q.partialNamespace = cmp.Or(q.partialNamespace, "partial")
	if q.evaluated != nil && q.compiler != nil {
		q.evaluated.WithAnnotationSet(q.compiler.GetAnnotationSet())
	}
	if q.seed == nil {
		q.seed = rand.Reader
	}
	if q.time.IsZero() {
		q.time = time.Now()
	}
	q.metrics = util.Or(q.metrics, metrics.New)

	f := &queryIDFactory{}
	b := newBindings(q.instr)

	e := &eval{
		ctx:                         ctx,
		metrics:                     q.metrics,
		seed:                        q.seed,
		timeStart:                   q.time.UnixNano(),
		cancel:                      q.cancel,
		query:                       q.query,
		queryCompiler:               q.queryCompiler,
		queryIDFact:                 f,
		queryID:                     f.Next(),
		bindings:                    b,
		compiler:                    q.compiler,
		store:                       q.store,
		baseCache:                   util.Or(q.baseCache, newBaseCache),
		txn:                         q.txn,
		input:                       q.input,
		external:                    q.external,
		tracers:                     q.tracers,
		traceEnabled:                len(q.tracers) > 0,
		stackCapture:                q.newStackCapture(),
		plugTraceVars:               q.plugTraceVars,
		instr:                       q.instr,
		builtins:                    q.builtins,
		builtinCache:                builtins.Cache{},
		interQueryBuiltinCache:      q.interQueryBuiltinCache,
		interQueryBuiltinValueCache: q.interQueryBuiltinValueCache,
		ndBuiltinCache:              q.ndBuiltinCache,
		virtualCache:                util.Or(q.virtualCache, NewVirtualCache),
		saveSet:                     newSaveSet(q.unknowns, b, q.instr),
		saveStack:                   newSaveStack(),
		saveSupport:                 newSaveSupport(),
		saveNamespace:               ast.InternedTerm(q.partialNamespace),
		skipSaveNamespace:           q.skipSaveNamespace,
		inliningControl: &inliningControl{
			shallow:                  q.shallowInlining,
			nondeterministicBuiltins: q.nondeterministicBuiltins,
		},
		genvarprefix:     q.genvarprefix,
		runtime:          q.runtime,
		indexing:         q.indexing,
		earlyExit:        q.earlyExit,
		builtinErrors:    &builtinErrors{},
		printHook:        q.printHook,
		strictObjects:    q.strictObjects,
		requestMetadata:  q.requestMetadata,
		responseMetadata: q.responseMetadata,
		evaluated:        q.evaluated,
	}

	if len(q.disableInlining) > 0 {
		e.inliningControl.PushDisable(q.disableInlining, false)
	}

	e.caller = e
	q.metrics.Timer(metrics.RegoPartialEval).Start()
	defer q.metrics.Timer(metrics.RegoPartialEval).Stop()

	livevars := ast.NewVarSet()
	for _, t := range q.unknowns {
		switch v := t.Value.(type) {
		case ast.Var:
			livevars.Add(v)
		case ast.Ref:
			livevars.Add(v[0].Value.(ast.Var))
		}
	}

	// iterate expressions to avoid Body -> any boxing in WalkVars argument
	for _, expr := range q.query {
		ast.WalkVars(expr, func(x ast.Var) bool {
			if !x.IsGenerated() {
				livevars.Add(x)
			}
			return false
		})
	}

	p := copypropagation.New(livevars).WithCompiler(q.compiler)

	err = e.Run(func(e *eval) error {
		// Build output from saved expressions.
		saved := e.saveStack.Peek()
		exprs := make([]*ast.Expr, 0, len(saved)+e.bindings.size())
		for _, elem := range saved {
			exprs = append(exprs, elem.Plug(e.bindings))
		}

		// Include bindings as exprs so that when caller evals the result, they
		// can obtain values for the vars in their query.
		_ = e.bindings.Iter(e.bindings, func(a, b *ast.Term) error {
			exprs = append(exprs, ast.Equality.Expr(a, b))
			return nil
		}) // cannot return error

		// Sort binding expressions so that results are deterministic.
		slices.SortFunc(exprs[len(saved):], (*ast.Expr).Compare)

		body := ast.NewBody(exprs...)

		// Skip this rule body if it fails to type-check.
		// Type-checking failure means the rule body will never succeed.
		if !e.compiler.PassesTypeCheck(body) {
			return nil
		}

		if !q.shallowInlining {
			body = applyCopyPropagation(p, e.instr, body)
		}

		partials = append(partials, body)
		return nil
	})

	if len(e.builtinErrors.errs) > 0 {
		if q.strictBuiltinErrors {
			err = e.builtinErrors.errs[0]
		} else if q.builtinErrorList != nil {
			// If a builtinErrorList has been supplied, we must use pointer indirection
			// to append to it. builtinErrorList is a slice pointer so that errors can be
			// appended to it without returning a new slice and changing the interface
			// of PartialRun.
			for _, err := range e.builtinErrors.errs {
				if tdError, ok := err.(*Error); ok {
					*(q.builtinErrorList) = append(*(q.builtinErrorList), *tdError)
				} else {
					*(q.builtinErrorList) = append(*(q.builtinErrorList), Error{
						Code:    BuiltinErr,
						Message: err.Error(),
					})
				}
			}
		}
	}

	support = e.saveSupport.List()

	for i, m := range support {
		if regoVersion := q.compiler.DefaultRegoVersion(); regoVersion != ast.RegoUndefined {
			ast.SetModuleRegoVersion(m, q.compiler.DefaultRegoVersion())
		}
		slices.SortFunc(support[i].Rules, (*ast.Rule).Compare)
	}

	return partials, support, err
}

// Run is a wrapper around Iter that accumulates query results and returns them
// in one shot.
func (q *Query) Run(ctx context.Context) (QueryResultSet, error) {
	qrs := QueryResultSet{}
	return qrs, q.Iter(ctx, func(qr QueryResult) error {
		qrs = append(qrs, qr)
		return nil
	})
}

// Iter executes the query and invokes the iter function with query results
// produced by evaluating the query.
func (q *Query) Iter(ctx context.Context, iter func(QueryResult) error) error {
	// Query evaluation must not be allowed if the compiler has errors and is in an undefined, possibly inconsistent state
	if q.compiler != nil && len(q.compiler.Errors) > 0 {
		return &Error{Code: InternalErr, Message: "compiler has errors"}
	}
	if q.evaluated != nil && q.compiler != nil {
		q.evaluated.WithAnnotationSet(q.compiler.GetAnnotationSet())
	}
	if q.seed == nil {
		q.seed = rand.Reader
	}
	if q.time.IsZero() {
		q.time = time.Now()
	}
	q.metrics = util.Or(q.metrics, metrics.New)

	f := &queryIDFactory{}
	e := &eval{
		ctx:                         ctx,
		metrics:                     q.metrics,
		seed:                        q.seed,
		timeStart:                   q.time.UnixNano(),
		cancel:                      q.cancel,
		query:                       q.query,
		queryCompiler:               q.queryCompiler,
		queryIDFact:                 f,
		queryID:                     f.Next(),
		bindings:                    newBindings(q.instr),
		compiler:                    q.compiler,
		store:                       q.store,
		baseCache:                   util.Or(q.baseCache, newBaseCache),
		txn:                         q.txn,
		input:                       q.input,
		external:                    q.external,
		tracers:                     q.tracers,
		traceEnabled:                len(q.tracers) > 0,
		stackCapture:                q.newStackCapture(),
		plugTraceVars:               q.plugTraceVars,
		instr:                       q.instr,
		builtins:                    q.builtins,
		builtinCache:                builtins.Cache{},
		interQueryBuiltinCache:      q.interQueryBuiltinCache,
		interQueryBuiltinValueCache: q.interQueryBuiltinValueCache,
		ndBuiltinCache:              q.ndBuiltinCache,
		virtualCache:                util.Or(q.virtualCache, NewVirtualCache),
		genvarprefix:                q.genvarprefix,
		runtime:                     q.runtime,
		indexing:                    q.indexing,
		earlyExit:                   q.earlyExit,
		builtinErrors:               &builtinErrors{},
		printHook:                   q.printHook,
		tracingOpts:                 q.tracingOpts,
		strictObjects:               q.strictObjects,
		roundTripper:                q.roundTripper,
		requestMetadata:             q.requestMetadata,
		responseMetadata:            q.responseMetadata,
		evaluated:                   q.evaluated,
	}
	if e.requestMetadata == nil {
		e.requestMetadata = map[string]any{}
	}
	e.caller = e
	q.metrics.Timer(metrics.RegoQueryEval).Start()
	err := e.Run(func(e *eval) error {
		qr := make(QueryResult, e.bindings.size())
		_ = e.bindings.Iter(nil, func(k, v *ast.Term) error {
			qr[k.Value.(ast.Var)] = v
			return nil
		}) // cannot return error
		return iter(qr)
	})

	if len(e.builtinErrors.errs) > 0 {
		if q.strictBuiltinErrors {
			err = e.builtinErrors.errs[0]
		} else if q.builtinErrorList != nil {
			// If a builtinErrorList has been supplied, we must use pointer indirection
			// to append to it. builtinErrorList is a slice pointer so that errors can be
			// appended to it without returning a new slice and changing the interface
			// of Iter.
			for _, err := range e.builtinErrors.errs {
				if tdError, ok := err.(*Error); ok {
					*(q.builtinErrorList) = append(*(q.builtinErrorList), *tdError)
				} else {
					*(q.builtinErrorList) = append(*(q.builtinErrorList), Error{
						Code:    BuiltinErr,
						Message: err.Error(),
					})
				}
			}
		}
	}

	q.metrics.Timer(metrics.RegoQueryEval).Stop()
	return err
}