// Copyright 2018 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 (
	"github.com/open-policy-agent/opa/v1/ast"
	"github.com/open-policy-agent/opa/v1/topdown/builtins"
)

func builtinArrayConcat(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	arrA, err := builtins.ArrayOperand(operands[0].Value, 1)
	if err != nil {
		return err
	}

	arrB, err := builtins.ArrayOperand(operands[1].Value, 2)
	if err != nil {
		return err
	}

	if arrA.Len() == 0 {
		return iter(operands[1])
	}
	if arrB.Len() == 0 {
		return iter(operands[0])
	}

	arrC := make([]*ast.Term, arrA.Len()+arrB.Len())

	i := 0
	arrA.Foreach(func(elemA *ast.Term) {
		arrC[i] = elemA
		i++
	})

	arrB.Foreach(func(elemB *ast.Term) {
		arrC[i] = elemB
		i++
	})

	return iter(ast.ArrayTerm(arrC...))
}

func builtinArrayFlatten(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	arr, err := builtins.ArrayOperand(operands[0].Value, 1)
	if err != nil {
		return err
	}

	size := arr.Len()
	preAlloc := size
	containsArray := false

	for i := range size {
		if nested, ok := arr.Elem(i).Value.(*ast.Array); ok {
			containsArray = true
			preAlloc += nested.Len() - 1
		}
	}

	if !containsArray && size == preAlloc {
		return iter(operands[0]) // Empty array, or no nested arrays -> nothing to flatten.
	}

	flattened := make([]*ast.Term, 0, preAlloc)
	for i := range size {
		elem := arr.Elem(i)
		if nested, ok := elem.Value.(*ast.Array); ok {
			for j := range nested.Len() {
				flattened = append(flattened, nested.Elem(j))
			}
		} else {
			flattened = append(flattened, elem)
		}
	}

	return iter(ast.ArrayTerm(flattened...))
}

func builtinArraySlice(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	arr, err := builtins.ArrayOperand(operands[0].Value, 1)
	if err != nil {
		return err
	}

	startIndex, err := builtins.IntOperand(operands[1].Value, 2)
	if err != nil {
		return err
	}

	stopIndex, err := builtins.IntOperand(operands[2].Value, 3)
	if err != nil {
		return err
	}

	l := arr.Len()

	// Clamp stopIndex to avoid out-of-range errors. If negative, clamp to zero.
	// Otherwise, clamp to length of array.
	if stopIndex < 0 {
		stopIndex = 0
	} else if stopIndex > l {
		stopIndex = l
	}

	// Clamp startIndex to avoid out-of-range errors. If negative, clamp to zero.
	// Otherwise, clamp to stopIndex to avoid to avoid cases like arr[1:0].
	if startIndex < 0 {
		startIndex = 0
	} else if startIndex > stopIndex {
		startIndex = stopIndex
	}

	if startIndex == 0 && stopIndex >= l {
		return iter(operands[0])
	}

	return iter(ast.NewTerm(arr.Slice(startIndex, stopIndex)))
}

func builtinArrayReverse(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	arr, err := builtins.ArrayOperand(operands[0].Value, 1)
	if err != nil {
		return err
	}

	length := arr.Len()

	if length == 0 {
		return iter(ast.InternedEmptyArray)
	}

	if length == 1 {
		return iter(operands[0])
	}

	reversedArr := make([]*ast.Term, length)
	for index := range length {
		reversedArr[index] = arr.Elem(length - index - 1)
	}

	return iter(ast.ArrayTerm(reversedArr...))
}

func init() {
	RegisterBuiltinFunc(ast.ArrayConcat.Name, builtinArrayConcat)
	RegisterBuiltinFunc(ast.ArrayFlatten.Name, builtinArrayFlatten)
	RegisterBuiltinFunc(ast.ArraySlice.Name, builtinArraySlice)
	RegisterBuiltinFunc(ast.ArrayReverse.Name, builtinArrayReverse)
}