// 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 (
	"errors"
	"math/big"

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

type arithArity1 func(a *big.Float) (*big.Float, error)
type arithArity2 func(a, b *big.Float) (*big.Float, error)

func arithAbs(a *big.Float) (*big.Float, error) {
	return a.Abs(a), nil
}

var halfAwayFromZero = big.NewFloat(0.5)

func arithRound(a *big.Float) (*big.Float, error) {
	var i *big.Int
	if a.Signbit() {
		i, _ = new(big.Float).Sub(a, halfAwayFromZero).Int(nil)
	} else {
		i, _ = new(big.Float).Add(a, halfAwayFromZero).Int(nil)
	}
	return new(big.Float).SetInt(i), nil
}

func arithCeil(a *big.Float) (*big.Float, error) {
	i, _ := a.Int(nil)
	f := new(big.Float).SetInt(i)

	if f.Signbit() || a.Cmp(f) == 0 {
		return f, nil
	}

	return new(big.Float).Add(f, big.NewFloat(1.0)), nil
}

func arithFloor(a *big.Float) (*big.Float, error) {
	i, _ := a.Int(nil)
	f := new(big.Float).SetInt(i)

	if !f.Signbit() || a.Cmp(f) == 0 {
		return f, nil
	}

	return new(big.Float).Sub(f, big.NewFloat(1.0)), nil
}

// exactIntArith applies op to n1 and n2 as exact big.Ints when both are integers.
//
// The big.Float path used otherwise carries the default mantissa, so integers needing more
// significant bits than that are silently rounded before the operation is applied.
func exactIntArith(n1, n2 ast.Number, op func(z, x, y *big.Int) *big.Int) (ast.Number, bool) {
	x, err := builtins.NumberToInt(n1)
	if err != nil {
		return "", false
	}
	y, err := builtins.NumberToInt(n2)
	if err != nil {
		return "", false
	}
	return builtins.IntToNumber(op(new(big.Int), x, y)), true
}

func builtinPlus(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	n1, err := builtins.NumberOperand(operands[0].Value, 1)
	if err != nil {
		return err
	}
	n2, err := builtins.NumberOperand(operands[1].Value, 2)
	if err != nil {
		return err
	}

	x, ok1 := n1.Int()
	y, ok2 := n2.Int()

	if ok1 && ok2 && inSmallIntRange(x) && inSmallIntRange(y) {
		return iter(ast.InternedTerm(x + y))
	}

	if n, ok := exactIntArith(n1, n2, (*big.Int).Add); ok {
		return iter(ast.NewTerm(n))
	}

	f := new(big.Float).Add(builtins.NumberToFloat(n1), builtins.NumberToFloat(n2))

	return iter(ast.NewTerm(builtins.FloatToNumber(f)))
}

func builtinMultiply(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	n1, err := builtins.NumberOperand(operands[0].Value, 1)
	if err != nil {
		return err
	}
	n2, err := builtins.NumberOperand(operands[1].Value, 2)
	if err != nil {
		return err
	}

	x, ok1 := n1.Int()
	y, ok2 := n2.Int()

	if ok1 && ok2 && inSmallIntRange(x) && inSmallIntRange(y) {
		return iter(ast.InternedTerm(x * y))
	}

	if n, ok := exactIntArith(n1, n2, (*big.Int).Mul); ok {
		return iter(ast.NewTerm(n))
	}

	f := new(big.Float).Mul(builtins.NumberToFloat(n1), builtins.NumberToFloat(n2))

	return iter(ast.NewTerm(builtins.FloatToNumber(f)))
}

func arithDivide(a, b *big.Float) (*big.Float, error) {
	i, acc := b.Int64()
	if acc == big.Exact && i == 0 {
		return nil, errors.New("divide by zero")
	}
	return new(big.Float).Quo(a, b), nil
}

func arithRem(a, b *big.Int) (*big.Int, error) {
	// Sign, not Int64: Int64 returns the low 64 bits when b does not fit in an
	// int64, so any nonzero multiple of 2^64 (e.g. 10 % 18446744073709551616)
	// would be misreported as modulo by zero.
	if b.Sign() == 0 {
		return nil, errors.New("modulo by zero")
	}
	return new(big.Int).Rem(a, b), nil
}

func builtinArithArity1(fn arithArity1) BuiltinFunc {
	return func(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
		n, err := builtins.NumberOperand(operands[0].Value, 1)
		if err != nil {
			return err
		}
		f, err := fn(builtins.NumberToFloat(n))
		if err != nil {
			return err
		}
		return iter(ast.NewTerm(builtins.FloatToNumber(f)))
	}
}

func builtinArithArity2(fn arithArity2) BuiltinFunc {
	return func(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
		n1, err := builtins.NumberOperand(operands[0].Value, 1)
		if err != nil {
			return err
		}
		n2, err := builtins.NumberOperand(operands[1].Value, 2)
		if err != nil {
			return err
		}
		f, err := fn(builtins.NumberToFloat(n1), builtins.NumberToFloat(n2))
		if err != nil {
			return err
		}
		return iter(ast.NewTerm(builtins.FloatToNumber(f)))
	}
}

func builtinMinus(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {

	n1, ok1 := operands[0].Value.(ast.Number)
	n2, ok2 := operands[1].Value.(ast.Number)

	if ok1 && ok2 {

		x, okx := n1.Int()
		y, oky := n2.Int()

		if okx && oky && inSmallIntRange(x) && inSmallIntRange(y) {
			return iter(ast.InternedTerm(x - y))
		}

		if n, ok := exactIntArith(n1, n2, (*big.Int).Sub); ok {
			return iter(ast.NewTerm(n))
		}

		f := new(big.Float).Sub(builtins.NumberToFloat(n1), builtins.NumberToFloat(n2))

		return iter(ast.NewTerm(builtins.FloatToNumber(f)))
	}

	s1, ok3 := operands[0].Value.(ast.Set)
	s2, ok4 := operands[1].Value.(ast.Set)

	if ok3 && ok4 {
		diff := s1.Diff(s2)
		if diff.Len() == 0 {
			return iter(ast.InternedEmptySet)
		}
		return iter(ast.NewTerm(diff))
	}

	if !ok1 && !ok3 {
		return builtins.NewOperandTypeErr(1, operands[0].Value, "number", "set")
	}

	if ok2 {
		return builtins.NewOperandTypeErr(2, operands[1].Value, "set")
	}

	return builtins.NewOperandTypeErr(2, operands[1].Value, "number")
}

func builtinRem(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	n1, ok1 := operands[0].Value.(ast.Number)
	n2, ok2 := operands[1].Value.(ast.Number)

	if ok1 && ok2 {

		x, okx := n1.Int()
		y, oky := n2.Int()

		if okx && oky && inSmallIntRange(x) && inSmallIntRange(y) {
			if y == 0 {
				return errors.New("modulo by zero")
			}

			return iter(ast.InternedTerm(x % y))
		}

		op1, err1 := builtins.NumberToInt(n1)
		op2, err2 := builtins.NumberToInt(n2)

		if err1 != nil || err2 != nil {
			return errors.New("modulo on floating-point number")
		}

		i, err := arithRem(op1, op2)
		if err != nil {
			return err
		}
		return iter(ast.NewTerm(builtins.IntToNumber(i)))
	}

	if !ok1 {
		return builtins.NewOperandTypeErr(1, operands[0].Value, "number")
	}

	return builtins.NewOperandTypeErr(2, operands[1].Value, "number")
}

func inSmallIntRange(num int) bool {
	return -1000 < num && num < 1000
}

func init() {
	RegisterBuiltinFunc(ast.Abs.Name, builtinArithArity1(arithAbs))
	RegisterBuiltinFunc(ast.Round.Name, builtinArithArity1(arithRound))
	RegisterBuiltinFunc(ast.Ceil.Name, builtinArithArity1(arithCeil))
	RegisterBuiltinFunc(ast.Floor.Name, builtinArithArity1(arithFloor))
	RegisterBuiltinFunc(ast.Plus.Name, builtinPlus)
	RegisterBuiltinFunc(ast.Minus.Name, builtinMinus)
	RegisterBuiltinFunc(ast.Multiply.Name, builtinMultiply)
	RegisterBuiltinFunc(ast.Divide.Name, builtinArithArity2(arithDivide))
	RegisterBuiltinFunc(ast.Rem.Name, builtinRem)
}