// Package calc evaluates the arithmetic people send the bot: decimal
// numbers, + - * / % ^, signs and parentheses.
//
// The expression is read by a small parser of its own, because Go's
// grammar has no power operator (^ is XOR there), and computed by
// go/constant, which does exact rational arithmetic: 5 * 5/2 is exactly
// 12.5, and 0.1 + 0.2 is exactly 0.3, so a result carries no binary
// floating point noise until the moment it is formatted. A power is the
// exception: one with a fractional exponent, or too large to compute
// exactly, is computed in float64.
package calc

import (
	"errors"
	"go/constant"
	"go/token"
	"math"
	"math/big"
	"regexp"
	"strconv"
	"strings"
)

// MaxInputLength caps an expression, in bytes, and maxExactPowerBits
// caps a power, so that a message cannot make the bot do unbounded work.
const MaxInputLength = 256

// maxExactPowerBits caps a power computed exactly: its numerator and
// denominator together have at most this many bits, estimated before
// multiplying as the exponent times the bits in the base's numerator and
// denominator. A larger power is computed in float64, whose cost does not
// grow with it.
const maxExactPowerBits = 4096

// Results of magnitude plainUpper or more are written in exponent form
// (1e+21 rather than twenty-two digits), and so are fractions smaller
// than plainLower (1e-07 rather than 0.0000001).
const (
	plainUpper = 1e21
	plainLower = 1e-6
)

// The precedence of the binary operators: the higher, the tighter the
// operator binds.
const (
	sumPrecedence = iota + 1
	productPrecedence
	powerPrecedence
)

// Errors returned by Evaluate. The bot turns each into a reply.
var (
	ErrTooLong        = errors.New("expression too long")
	ErrNotArithmetic  = errors.New("not an arithmetic expression")
	ErrDivisionByZero = errors.New("division by zero")
	ErrTooLarge       = errors.New("result too large")
	ErrNoRealResult   = errors.New("no real result")
)

// decimal is the only number syntax accepted. Go's own literal syntax is
// wider, and parts of it are traps for someone typing arithmetic: 010 is
// octal 8, and 0x10, 1_000 and 1i are not what a calculator user means
// by a number. Here the x, _ or i matches no token and is refused.
const decimal = `([0-9]+\.?[0-9]*|\.[0-9]+)([eE][+-]?[0-9]+)?`

var (
	// nextToken matches the token at the start of the input, after any
	// whitespace: an operator, a parenthesis or a number. ** comes
	// before * so that it is read as one token.
	nextToken = regexp.MustCompile(`^\s*(\*\*|[-+*/%^()]|` + decimal + `)`)

	// decimalLiteral matches a token that is a number.
	decimalLiteral = regexp.MustCompile(`^` + decimal + `$`)
)

// Evaluate computes an arithmetic expression and returns its result as
// text: whole numbers without a decimal point, fractions in the
// shortest form that reads back as the same float64.
func Evaluate(input string) (string, error) {
	s := strings.TrimSpace(input)
	if len(s) > MaxInputLength {
		return "", ErrTooLong
	}

	tokens, err := tokenize(s)
	if err != nil {
		return "", err
	}

	p := parser{tokens: tokens}

	v, err := p.expr(sumPrecedence)
	if err != nil {
		return "", err
	}

	if p.next() != "" {
		return "", ErrNotArithmetic
	}

	return format(v)
}

// tokenize splits an expression into operators, parentheses and
// numbers, and refuses anything else. ** is returned as ^.
func tokenize(s string) ([]string, error) {
	var tokens []string

	for strings.TrimSpace(s) != "" {
		m := nextToken.FindStringSubmatch(s)
		if m == nil {
			return nil, ErrNotArithmetic
		}

		tok := m[1]
		if tok == "**" {
			tok = "^"
		}

		tokens = append(tokens, tok)
		s = s[len(m[0]):]
	}

	return tokens, nil
}

// parser computes an expression as it reads it, by precedence climbing:
// expr reads operands joined by operators of at least a given
// precedence, and hands the right operand of each to a deeper call that
// takes only the operators that bind tighter, so those are applied
// first.
type parser struct {
	tokens []string
}

// next removes and returns the next token, or "" at the end.
func (p *parser) next() string {
	tok := p.peek()
	if tok != "" {
		p.tokens = p.tokens[1:]
	}

	return tok
}

// peek returns the next token, or "" at the end, and leaves it unread.
func (p *parser) peek() string {
	if len(p.tokens) == 0 {
		return ""
	}

	return p.tokens[0]
}

// expr reads and computes an expression whose binary operators all have
// at least minPrecedence. Operators of equal precedence group to the
// left, 8/2/2 is (8/2)/2, except ^, which groups to the right: 2^3^2 is
// 2^(3^2).
func (p *parser) expr(minPrecedence int) (constant.Value, error) {
	x, err := p.operand()
	if err != nil {
		return nil, err
	}

	for {
		op := p.peek()

		prec := precedence(op)
		if prec < minPrecedence {
			return x, nil
		}

		p.next()

		rightPrecedence := prec + 1
		if op == "^" {
			rightPrecedence = prec
		}

		y, err := p.expr(rightPrecedence)
		if err != nil {
			return nil, err
		}

		x, err = apply(x, op, y)
		if err != nil {
			return nil, err
		}
	}
}

// operand reads a number, an expression in parentheses, or a sign and
// its operand. A sign binds more loosely than a power that follows it,
// so -2^2 is -(2^2), and 2^-1 is 2^(-1).
func (p *parser) operand() (constant.Value, error) {
	switch tok := p.next(); tok {
	case "+", "-":
		x, err := p.expr(powerPrecedence)
		if err != nil {
			return nil, err
		}

		if tok == "-" {
			x = constant.UnaryOp(token.SUB, x, 0)
		}

		return x, nil
	case "(":
		x, err := p.expr(sumPrecedence)
		if err != nil {
			return nil, err
		}

		if p.next() != ")" {
			return nil, ErrNotArithmetic
		}

		return x, nil
	default:
		return number(tok)
	}
}

// precedence returns the precedence of a binary operator, and 0 for any
// other token, which ends an expression.
func precedence(op string) int {
	switch op {
	case "+", "-":
		return sumPrecedence
	case "*", "/", "%":
		return productPrecedence
	case "^":
		return powerPrecedence
	default:
		return 0
	}
}

func number(tok string) (constant.Value, error) {
	if !decimalLiteral.MatchString(tok) {
		return nil, ErrNotArithmetic
	}

	// Read as FLOAT, which makes every literal decimal: as INT, a
	// leading zero would make it octal.
	v := constant.MakeFromLiteral(tok, token.FLOAT, 0)

	// The syntax was checked above, so Unknown here means the exponent
	// overflowed.
	if v.Kind() == constant.Unknown {
		return nil, ErrTooLarge
	}

	return v, nil
}

// apply computes x op y.
func apply(x constant.Value, op string, y constant.Value) (constant.Value, error) {
	var (
		v   constant.Value
		err error
	)

	switch op {
	case "+":
		v = constant.BinaryOp(x, token.ADD, y)
	case "-":
		v = constant.BinaryOp(x, token.SUB, y)
	case "*":
		v = constant.BinaryOp(x, token.MUL, y)
	case "/":
		v, err = divide(x, y)
	case "%":
		v, err = modulo(x, y)
	case "^":
		v, err = power(x, y)
	default:
		err = ErrNotArithmetic
	}

	if err != nil {
		return nil, err
	}

	// go/constant represents an overflow to infinity as Unknown.
	if v.Kind() == constant.Unknown {
		return nil, ErrTooLarge
	}

	return v, nil
}

func divide(x, y constant.Value) (constant.Value, error) {
	// constant.BinaryOp panics on a zero divisor.
	if constant.Sign(y) == 0 {
		return nil, ErrDivisionByZero
	}

	// token.QUO divides exactly, integers included: 25/2 is 12.5.
	return constant.BinaryOp(x, token.QUO, y), nil
}

// modulo computes x % y, whose result takes the sign of y, as in Python:
// -7 % 3 is 2 and 7 % -3 is -2. It is exact for decimals too: 7.5 % 2
// is 1.5.
func modulo(x, y constant.Value) (constant.Value, error) {
	q, err := divide(x, y)
	if err != nil {
		return nil, err
	}

	// The whole part of a rounded quotient, and so the remainder, would
	// be wrong.
	if rounded(q) {
		return nil, ErrTooLarge
	}

	// token.QUO_ASSIGN is go/constant's integer division, which
	// truncates, so r has the sign of x and is less than y in size.
	whole := constant.BinaryOp(constant.Num(q), token.QUO_ASSIGN, constant.Denom(q))
	r := constant.BinaryOp(x, token.SUB, constant.BinaryOp(y, token.MUL, whole))

	if constant.Sign(r) != 0 && constant.Sign(r) != constant.Sign(y) {
		r = constant.BinaryOp(r, token.ADD, y)
	}

	return r, nil
}

// power computes x^y. A negative x needs a whole y, and its sign is
// applied here: math.Pow would take it from the parity of y's float64
// value, and every float64 from 2^53 up is even.
func power(x, y constant.Value) (constant.Value, error) {
	// n is y if y is a whole number, and Unknown otherwise.
	n := constant.ToInt(y)

	switch {
	case constant.Sign(x) == 0 && constant.Sign(y) < 0:
		return nil, ErrDivisionByZero
	case constant.Sign(x) >= 0:
		return nonNegativePower(x, y, n), nil
	case rounded(y):
		// Whether a rounded y is whole, or odd, is unknown.
		return nil, ErrTooLarge
	case n.Kind() != constant.Int:
		return nil, ErrNoRealResult
	}

	// x is negative and n whole: x^n is (-x)^n, negated if n is odd.
	v := nonNegativePower(constant.UnaryOp(token.SUB, x, 0), y, n)

	odd := constant.BinaryOp(n, token.AND, constant.MakeInt64(1))
	if constant.Sign(odd) != 0 {
		v = constant.UnaryOp(token.SUB, v, 0)
	}

	return v, nil
}

// nonNegativePower computes x^y for x of at least zero, and y not below
// zero if x is zero: exactly if y is a whole number n and the result
// fits in maxExactPowerBits, otherwise in float64.
func nonNegativePower(x, y, n constant.Value) constant.Value {
	e, ok := constant.Int64Val(n)
	if ok && exactPowerFits(x, e) {
		return exactPower(x, e)
	}

	xf, _ := constant.Float64Val(x)
	yf, _ := constant.Float64Val(y)

	// An infinite result becomes Unknown, which apply refuses as too
	// large.
	return constant.MakeFloat64(math.Pow(xf, yf))
}

// exactPowerFits reports whether x^e fits in maxExactPowerBits.
func exactPowerFits(x constant.Value, e int64) bool {
	// Checked first so that the product below cannot overflow.
	if e < -maxExactPowerBits || e > maxExactPowerBits {
		return false
	}

	// Num and Denom are Unknown for a value too large or too small to
	// be held as a fraction.
	num, den := constant.Num(x), constant.Denom(x)
	if num.Kind() != constant.Int {
		return false
	}

	bits := int64(constant.BitLen(num) + constant.BitLen(den))

	return bits*max(e, -e) <= maxExactPowerBits
}

// exactPower computes x^e by repeated squaring. x is not zero if e is
// negative.
func exactPower(x constant.Value, e int64) constant.Value {
	result := constant.MakeInt64(1)

	for n := max(e, -e); n > 0; n >>= 1 {
		if n&1 == 1 {
			result = constant.BinaryOp(result, token.MUL, x)
		}

		x = constant.BinaryOp(x, token.MUL, x)
	}

	if e < 0 {
		result = constant.BinaryOp(constant.MakeInt64(1), token.QUO, result)
	}

	return result
}

// rounded reports whether go/constant holds v rounded. It holds a
// number exactly, as a fraction, only while the numerator and the
// denominator each stay under 4096 bits; past that, and for a literal of
// that size, it holds a 512-bit float.
func rounded(v constant.Value) bool {
	_, isFloat := constant.Val(v).(*big.Float)

	return isFloat
}

// format writes a result for a person to read. A whole number of
// ordinary size is written exactly, digit for digit; anything else goes
// through float64, whose shortest round-trip form is free of the noise
// (0.30000000000000004) that printing a binary fraction to a fixed
// precision produces.
func format(v constant.Value) (string, error) {
	f, _ := constant.Float64Val(v)
	if math.IsInf(f, 0) || math.IsNaN(f) {
		return "", ErrTooLarge
	}

	abs := math.Abs(f)

	if i := constant.ToInt(v); i.Kind() == constant.Int && abs < plainUpper {
		return i.ExactString(), nil
	}

	if abs >= plainUpper || abs < plainLower {
		return strconv.FormatFloat(f, 'g', -1, 64), nil
	}

	return strconv.FormatFloat(f, 'f', -1, 64), nil
}