// Copyright 2016 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 topdown

import (
	"context"
	"encoding/binary"
	"io"
	"math/rand"

	"github.com/open-policy-agent/opa/v1/ast"
	"github.com/open-policy-agent/opa/v1/metrics"
	"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/tracing"
)

type (
	// Deprecated: Functional-style builtins are deprecated. Use BuiltinFunc instead.
	FunctionalBuiltin1 func(op1 ast.Value) (output ast.Value, err error)

	// Deprecated: Functional-style builtins are deprecated. Use BuiltinFunc instead.
	FunctionalBuiltin2 func(op1, op2 ast.Value) (output ast.Value, err error)

	// Deprecated: Functional-style builtins are deprecated. Use BuiltinFunc instead.
	FunctionalBuiltin3 func(op1, op2, op3 ast.Value) (output ast.Value, err error)

	// Deprecated: Functional-style builtins are deprecated. Use BuiltinFunc instead.
	FunctionalBuiltin4 func(op1, op2, op3, op4 ast.Value) (output ast.Value, err error)

	// BuiltinContext contains context from the evaluator that may be used by
	// built-in functions.
	BuiltinContext struct {
		Context                     context.Context            // request context that was passed when query started
		Metrics                     metrics.Metrics            // metrics registry for recording built-in specific metrics
		Seed                        io.Reader                  // randomization source
		Time                        *ast.Term                  // wall clock time
		Cancel                      Cancel                     // atomic value that signals evaluation to halt
		Runtime                     *ast.Term                  // runtime information on the OPA instance
		Cache                       builtins.Cache             // built-in function state cache
		InterQueryBuiltinCache      cache.InterQueryCache      // cross-query built-in function state cache
		InterQueryBuiltinValueCache cache.InterQueryValueCache // cross-query built-in function state value cache. this cache is useful for scenarios where the entry size cannot be calculated
		NDBuiltinCache              builtins.NDBCache          // cache for non-deterministic built-in state
		Location                    *ast.Location              // location of built-in call
		Tracers                     []Tracer                   // Deprecated: Use QueryTracers instead
		QueryTracers                []QueryTracer              // tracer objects for trace() built-in function
		TraceEnabled                bool                       // indicates whether tracing is enabled for the evaluation
		QueryID                     uint64                     // identifies query being evaluated
		ParentID                    uint64                     // identifies parent of query being evaluated
		PrintHook                   print.Hook                 // provides callback function to use for printing
		RoundTripper                CustomizeRoundTripper      // customize transport to use for HTTP requests
		DistributedTracingOpts      tracing.Options            // options to be used by distributed tracing.
		RequestMetadata             map[string]any             // metadata from the caller, for use by wrapping projects
		ResponseMetadata            map[string]any             // metadata for the response, populated by wrapping projects
		rand                        *rand.Rand                 // randomization source for non-security-sensitive operations
		Capabilities                *ast.Capabilities
	}

	// BuiltinFunc defines an interface for implementing built-in functions.
	// The built-in function is called with the plugged operands from the call
	// (including the output operands.) The implementation should evaluate the
	// operands and invoke the iterator for each successful/defined output
	// value.
	BuiltinFunc func(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error
)

// Rand returns a random number generator based on the Seed for this built-in
// context. The random number will be re-used across multiple calls to this
// function. If a random number generator cannot be created, an error is
// returned.
func (bctx *BuiltinContext) Rand() (*rand.Rand, error) {
	if bctx.rand != nil {
		return bctx.rand, nil
	}

	seed, err := readInt64(bctx.Seed)
	if err != nil {
		return nil, err
	}

	bctx.rand = rand.New(rand.NewSource(seed))
	return bctx.rand, nil
}

// RegisterBuiltinFunc adds a new built-in function to the evaluation engine.
func RegisterBuiltinFunc(name string, f BuiltinFunc) {
	builtinFunctions[name] = builtinErrorWrapper(name, f)
}

// Deprecated: Functional-style builtins are deprecated. Use RegisterBuiltinFunc instead.
func RegisterFunctionalBuiltin1(name string, fun FunctionalBuiltin1) {
	builtinFunctions[name] = functionalWrapper1(name, fun)
}

// Deprecated: Functional-style builtins are deprecated. Use RegisterBuiltinFunc instead.
func RegisterFunctionalBuiltin2(name string, fun FunctionalBuiltin2) {
	builtinFunctions[name] = functionalWrapper2(name, fun)
}

// Deprecated: Functional-style builtins are deprecated. Use RegisterBuiltinFunc instead.
func RegisterFunctionalBuiltin3(name string, fun FunctionalBuiltin3) {
	builtinFunctions[name] = functionalWrapper3(name, fun)
}

// Deprecated: Functional-style builtins are deprecated. Use RegisterBuiltinFunc instead.
func RegisterFunctionalBuiltin4(name string, fun FunctionalBuiltin4) {
	builtinFunctions[name] = functionalWrapper4(name, fun)
}

// GetBuiltin returns a built-in function implementation, nil if no built-in found.
func GetBuiltin(name string) BuiltinFunc {
	return builtinFunctions[name]
}

// Deprecated: The BuiltinEmpty type is no longer needed. Use nil return values instead.
type BuiltinEmpty struct{}

func (BuiltinEmpty) Error() string {
	return "<empty>"
}

var builtinFunctions = map[string]BuiltinFunc{}

func builtinErrorWrapper(name string, fn BuiltinFunc) BuiltinFunc {
	return func(bctx BuiltinContext, args []*ast.Term, iter func(*ast.Term) error) error {
		err := fn(bctx, args, iter)
		if err == nil {
			return nil
		}
		return handleBuiltinErr(name, bctx.Location, err)
	}
}

func functionalWrapper1(name string, fn FunctionalBuiltin1) BuiltinFunc {
	return func(bctx BuiltinContext, args []*ast.Term, iter func(*ast.Term) error) error {
		result, err := fn(args[0].Value)
		if err == nil {
			return iter(ast.NewTerm(result))
		}
		return handleBuiltinErr(name, bctx.Location, err)
	}
}

func functionalWrapper2(name string, fn FunctionalBuiltin2) BuiltinFunc {
	return func(bctx BuiltinContext, args []*ast.Term, iter func(*ast.Term) error) error {
		result, err := fn(args[0].Value, args[1].Value)
		if err == nil {
			return iter(ast.NewTerm(result))
		}
		return handleBuiltinErr(name, bctx.Location, err)
	}
}

func functionalWrapper3(name string, fn FunctionalBuiltin3) BuiltinFunc {
	return func(bctx BuiltinContext, args []*ast.Term, iter func(*ast.Term) error) error {
		result, err := fn(args[0].Value, args[1].Value, args[2].Value)
		if err == nil {
			return iter(ast.NewTerm(result))
		}
		return handleBuiltinErr(name, bctx.Location, err)
	}
}

func functionalWrapper4(name string, fn FunctionalBuiltin4) BuiltinFunc {
	return func(bctx BuiltinContext, args []*ast.Term, iter func(*ast.Term) error) error {
		result, err := fn(args[0].Value, args[1].Value, args[2].Value, args[3].Value)
		if err == nil {
			return iter(ast.NewTerm(result))
		}
		if _, empty := err.(BuiltinEmpty); empty {
			return nil
		}
		return handleBuiltinErr(name, bctx.Location, err)
	}
}

func handleBuiltinErr(name string, loc *ast.Location, err error) error {
	var code string
	switch err := err.(type) {
	case BuiltinEmpty:
		return nil
	case *Error, Halt:
		return err
	case builtins.ErrOperand:
		code = TypeErr
	default:
		code = BuiltinErr
	}
	return (&Error{Code: code, Message: name + ": " + err.Error(), Location: loc}).Wrap(err)
}

func readInt64(r io.Reader) (int64, error) {
	bs := make([]byte, 8)
	n, err := io.ReadFull(r, bs)
	if n != len(bs) || err != nil {
		return 0, err
	}
	return int64(binary.BigEndian.Uint64(bs)), nil
}

// Used to get older-style (ast.Term, error) tuples out of newer functions.
func getResult(fn BuiltinFunc, operands ...*ast.Term) (*ast.Term, error) {
	var result *ast.Term
	extractionFn := func(r *ast.Term) error {
		result = r
		return nil
	}
	err := fn(BuiltinContext{}, operands, extractionFn)
	if err != nil {
		return nil, err
	}
	return result, nil
}