// 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 builtins contains utilities for implementing built-in functions.
package builtins

import (
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"strings"

	"github.com/open-policy-agent/opa/v1/ast"
	"github.com/open-policy-agent/opa/v1/util"
)

// Cache defines the built-in cache used by the top-down evaluation. The keys
// must be comparable and should not be of type string.
type Cache map[any]any

// Put updates the cache for the named built-in.
func (c Cache) Put(k, v any) {
	c[k] = v
}

// Get returns the cached value for k.
func (c Cache) Get(k any) (any, bool) {
	v, ok := c[k]
	return v, ok
}

// We use an ast.Object for the cached keys/values because a naive
// map[ast.Value]ast.Value will not correctly detect value equality of
// the member keys.
type NDBCache map[string]ast.Object

func (c NDBCache) AsValue() ast.Value {
	out := ast.NewObject()
	for bname, obj := range c {
		out.Insert(ast.InternedTerm(bname), ast.NewTerm(obj))
	}
	return out
}

// Put updates the cache for the named built-in.
// Automatically creates the 2-level hierarchy as needed.
func (c NDBCache) Put(name string, k, v ast.Value) {
	if _, ok := c[name]; !ok {
		c[name] = ast.NewObject()
	}
	c[name].Insert(ast.NewTerm(k), ast.NewTerm(v))
}

// Get returns the cached value for k for the named builtin.
func (c NDBCache) Get(name string, k ast.Value) (ast.Value, bool) {
	if m, ok := c[name]; ok {
		v := m.Get(ast.NewTerm(k))
		if v != nil {
			return v.Value, true
		}
		return nil, false
	}
	return nil, false
}

// Convenience functions for serializing the data structure.
func (c NDBCache) MarshalJSON() ([]byte, error) {
	v, err := ast.JSON(c.AsValue())
	if err != nil {
		return nil, err
	}
	return json.Marshal(v)
}

func (c *NDBCache) UnmarshalJSON(data []byte) error {
	out := map[string]ast.Object{}
	var incoming any

	// Note: We use util.Unmarshal instead of json.Unmarshal to get
	// correct deserialization of number types.
	err := util.Unmarshal(data, &incoming)
	if err != nil {
		return err
	}

	// Convert interface types back into ast.Value types.
	nestedObject, err := ast.InterfaceToValue(incoming)
	if err != nil {
		return err
	}

	// Reconstruct NDBCache from nested ast.Object structure.
	if source, ok := nestedObject.(ast.Object); ok {
		err = source.Iter(func(k, v *ast.Term) error {
			if obj, ok := v.Value.(ast.Object); ok {
				out[string(k.Value.(ast.String))] = obj
				return nil
			}
			return errors.New("expected Object, got other Value type in conversion")
		})
		if err != nil {
			return err
		}
	}

	*c = out

	return nil
}

// ErrOperand represents an invalid operand has been passed to a built-in
// function. Built-ins should return ErrOperand to indicate a type error has
// occurred.
type ErrOperand string

func (err ErrOperand) Error() string {
	return string(err)
}

// NewOperandErr returns a generic operand error.
func NewOperandErr(pos int, f string, a ...any) error {
	f = fmt.Sprintf("operand %v ", pos) + f
	return ErrOperand(fmt.Sprintf(f, a...))
}

// NewOperandTypeErr returns an operand error indicating the operand's type was wrong.
func NewOperandTypeErr(pos int, got ast.Value, expected ...string) error {
	if len(expected) == 1 {
		return NewOperandErr(pos, "must be %v but got %v", expected[0], ast.ValueName(got))
	}

	return NewOperandErr(pos, "must be one of {%v} but got %v", strings.Join(expected, ", "), ast.ValueName(got))
}

// NewOperandElementErr returns an operand error indicating an element in the
// composite operand was wrong.
func NewOperandElementErr(pos int, composite ast.Value, got ast.Value, expected ...string) error {
	tpe := ast.ValueName(composite)

	if len(expected) == 1 {
		return NewOperandErr(pos, "must be %v of %vs but got %v containing %v", tpe, expected[0], tpe, ast.ValueName(got))
	}

	return NewOperandErr(pos, "must be %v of (any of) {%v} but got %v containing %v", tpe, strings.Join(expected, ", "), tpe, ast.ValueName(got))
}

// NewOperandEnumErr returns an operand error indicating a value was wrong.
func NewOperandEnumErr(pos int, expected ...string) error {
	if len(expected) == 1 {
		return NewOperandErr(pos, "must be %v", expected[0])
	}

	return NewOperandErr(pos, "must be one of {%v}", strings.Join(expected, ", "))
}

// IntOperand converts x to an int. If the cast fails, a descriptive error is
// returned.
func IntOperand(x ast.Value, pos int) (int, error) {
	n, ok := x.(ast.Number)
	if !ok {
		return 0, NewOperandTypeErr(pos, x, "number")
	}

	i, ok := n.Int()
	if !ok {
		return 0, NewOperandErr(pos, "must be integer number but got floating-point number")
	}

	return i, nil
}

// BigIntOperand converts x to a big int. If the cast fails, a descriptive error
// is returned.
func BigIntOperand(x ast.Value, pos int) (*big.Int, error) {
	n, err := NumberOperand(x, 1)
	if err != nil {
		return nil, NewOperandTypeErr(pos, x, "integer")
	}
	bi, err := NumberToInt(n)
	if err != nil {
		return nil, NewOperandErr(pos, "must be integer number but got floating-point number")
	}

	return bi, nil
}

// NumberOperand converts x to a number. If the cast fails, a descriptive error is
// returned.
func NumberOperand(x ast.Value, pos int) (ast.Number, error) {
	if n, ok := x.(ast.Number); ok {
		return n, nil
	}
	return ast.Number(""), NewOperandTypeErr(pos, x, "number")
}

// SetOperand converts x to a set. If the cast fails, a descriptive error is
// returned.
func SetOperand(x ast.Value, pos int) (ast.Set, error) {
	if s, ok := x.(ast.Set); ok {
		return s, nil
	}
	return nil, NewOperandTypeErr(pos, x, "set")
}

// StringOperand returns x as [ast.String], or a descriptive error if the conversion fails.
func StringOperand(x ast.Value, pos int) (ast.String, error) {
	if s, ok := x.(ast.String); ok {
		return s, nil
	}
	return ast.String(""), NewOperandTypeErr(pos, x, "string")
}

// StringOperandByteSlice returns x a []byte, assuming x is [ast.String], or a descriptive error
// if that is not the case. The returned byte slice points directly at the underlying array backing
// the string, and should not be modified.
func StringOperandByteSlice(x ast.Value, pos int) ([]byte, error) {
	s, err := StringOperand(x, pos)
	if err != nil {
		return nil, err
	}
	return util.StringToByteSlice(string(s)), nil
}

// ObjectOperand converts x to an object. If the conversion fails, a descriptive
// error is returned.
func ObjectOperand(x ast.Value, pos int) (o ast.Object, err error) {
	if o, ok := x.(ast.Object); ok {
		return o, nil
	}
	return nil, NewOperandTypeErr(pos, x, "object")
}

// ArrayOperand converts x to an array. If the conversion fails, a descriptive
// error is returned.
func ArrayOperand(x ast.Value, pos int) (*ast.Array, error) {
	if a, ok := x.(*ast.Array); ok {
		return a, nil
	}
	return nil, NewOperandTypeErr(pos, x, "array")
}

// NumberToFloat converts n to a big float.
func NumberToFloat(n ast.Number) *big.Float {
	return NumberToFloatInto(nil, n)
}

// NumberToFloatInto converts n to a big float, storing it in dst when provided.
func NumberToFloatInto(dst *big.Float, n ast.Number) *big.Float {
	if dst == nil {
		dst = new(big.Float)
	}
	if _, ok := dst.SetString(string(n)); !ok {
		panic("illegal value")
	}
	return dst
}

// FloatToNumber converts f to a number.
func FloatToNumber(f *big.Float) ast.Number {
	var format byte = 'g'
	if f.IsInt() {
		format = 'f'
	}
	return ast.Number(f.Text(format, -1))
}

// NumberToInt converts n to a big int.
// If n cannot be converted to an big int, an error is returned.
func NumberToInt(n ast.Number) (*big.Int, error) {
	// Integer literals are parsed exactly. Going through NumberToFloat first would round any
	// value needing more than the big.Float default mantissa, and because the rounded value is
	// itself an integer, the accuracy check below cannot detect that it happened.
	if i, ok := new(big.Int).SetString(string(n), 10); ok {
		return i, nil
	}
	// Fractional and exponent forms. big.Rat parses both exactly, so a value such as 1e30 stays
	// exact, and a genuinely fractional value is rejected rather than silently truncated.
	r, ok := new(big.Rat).SetString(string(n))
	if !ok || !r.IsInt() {
		return nil, errors.New("illegal value")
	}
	return new(big.Int).Set(r.Num()), nil
}

// IntToNumber converts i to a number.
func IntToNumber(i *big.Int) ast.Number {
	return ast.Number(i.String())
}

// StringSliceOperand converts x to a []string. If the cast fails, a descriptive error is
// returned.
func StringSliceOperand(a ast.Value, pos int) ([]string, error) {
	type iterable interface {
		Iter(func(*ast.Term) error) error
		Len() int
	}

	strs, ok := a.(iterable)
	if !ok {
		return nil, NewOperandTypeErr(pos, a, "array", "set")
	}

	var outStrs = make([]string, 0, strs.Len())
	if err := strs.Iter(func(x *ast.Term) error {
		s, ok := x.Value.(ast.String)
		if !ok {
			return NewOperandElementErr(pos, a, x.Value, "string")
		}
		outStrs = append(outStrs, string(s))
		return nil
	}); err != nil {
		return nil, err
	}

	return outStrs, nil
}

// RuneSliceOperand converts x to a []rune. If the cast fails, a descriptive error is
// returned.
func RuneSliceOperand(x ast.Value, pos int) ([]rune, error) {
	a, err := ArrayOperand(x, pos)
	if err != nil {
		return nil, err
	}

	var f = make([]rune, a.Len())
	for k := range a.Len() {
		b := a.Elem(k)
		c, ok := b.Value.(ast.String)
		if !ok {
			return nil, NewOperandElementErr(pos, x, b.Value, "string")
		}

		d := []rune(string(c))
		if len(d) != 1 {
			return nil, NewOperandElementErr(pos, x, b.Value, "rune")
		}

		f[k] = d[0]
	}

	return f, nil
}