package ast_test

import (
	"encoding"
	"testing"

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

var valueAppenderTests = []struct {
	name string
	term ast.StringLengther
	want string
}{
	{
		name: "var",
		term: ast.MustParseTerm("input.foo.bar"),
		want: "input.foo.bar",
	},
	{
		name: "string",
		term: ast.StringTerm(`foo bar baz`),
		want: `"foo bar baz"`,
	},
	{
		name: "string with escapes",
		term: ast.StringTerm(`"foo" "bar" "qux"`),
		want: `"\"foo\" \"bar\" \"qux\""`,
	},
	{
		name: "string with newlines",
		term: ast.StringTerm("line1\nline2\nline3"),
		want: `"line1\nline2\nline3"`,
	},
	{
		name: "number",
		term: ast.MustParseTerm("3.14"),
		want: "3.14",
	},
	{
		name: "boolean",
		term: ast.MustParseTerm("true"),
		want: "true",
	},
	{
		name: "null",
		term: ast.MustParseTerm("null"),
		want: "null",
	},
	{
		name: "array",
		term: ast.MustParseTerm(`[1, "two", false]`),
		want: `[1, "two", false]`,
	},
	{
		name: "set",
		term: ast.MustParseTerm(`{1, 2, 3}`),
		want: `{1, 2, 3}`,
	},
	{
		name: "object",
		term: ast.MustParseTerm(`{"a": 1, "b": "two"}`),
		want: `{"a": 1, "b": "two"}`,
	},
	{
		name: "call",
		term: ast.MustParseExpr(`foo(input.bar, "baz")`),
		want: `foo(input.bar, "baz")`,
	},
	{
		name: "template string",
		term: ast.MustParseTerm(`$"Hello, {input.name}!"`),
		want: `$"Hello, {input.name}!"`,
	},
	{
		name: "ref",
		term: ast.MustParseTerm(`data.foo["b a r"].allow`),
		want: `data.foo["b a r"].allow`,
	},
	{
		name: "object comprehension",
		term: ast.MustParseTerm(`{k: v |
				some k, v
				data.items[k] == v
				v > 10
			}`),
		want: `{k: v | some k, v; equal(data.items[k], v); gt(v, 10)}`,
	},
	{
		name: "object comprehension infix value",
		term: ast.MustParseTerm(`{k: (v * 100) | some k, v; v > 0}`),
		want: `{k: (v * 100) | some k, v; gt(v, 0)}`,
	},
	{
		name: "object comprehension infix key",
		term: ast.MustParseTerm(`{(k + 1): v | some k, v; v > 0}`),
		want: `{(k + 1): v | some k, v; gt(v, 0)}`,
	},
	{
		name: "array comprehension",
		term: ast.MustParseTerm(`[(x * 2) | x := data.numbers[_]; x > 5]`),
		want: `[(x * 2) | assign(x, data.numbers[_]); gt(x, 5)]`,
	},
	{
		name: "array comprehension nested infix operators",
		term: ast.MustParseTerm(`[((x + 1) * 2) | x := data.numbers[_]]`),
		want: `[((x + 1) * 2) | assign(x, data.numbers[_])]`,
	},
	{
		name: "array comprehension non-infix head",
		term: ast.MustParseTerm(`[count(x) | x := data.lists[_]]`),
		want: `[count(x) | assign(x, data.lists[_])]`,
	},
	{
		name: "array comprehension nested infix with function call",
		term: ast.MustParseTerm(`[(to_number(cpu_str[i])/1000) | cpu_str[i]]`),
		want: `[(to_number(cpu_str[i]) / 1000) | cpu_str[i]]`,
	},
	{
		name: "set comprehension",
		term: ast.MustParseTerm(`{x | x := data.values[_]; x < 100}`),
		want: `{x | assign(x, data.values[_]); lt(x, 100)}`,
	},
	{
		name: "not",
		term: &ast.Not{
			Body: ast.NewBody(ast.NewExpr(ast.LessThan.Call(ast.NumberTerm("100"), ast.NumberTerm("1")))),
		},
		want: `not lt(100, 1)`,
	},
	{
		name: "not, explicit body, one-line",
		term: &ast.Not{
			ExplicitBody: true,
			Body:         ast.NewBody(ast.NewExpr(ast.LessThan.Call(ast.NumberTerm("100"), ast.NumberTerm("1")))),
		},
		want: `not {lt(100, 1)}`,
	},
	{
		name: "not, implicit body, multi-line",
		term: &ast.Not{
			Body: ast.NewBody(
				ast.NewExpr(ast.LessThan.Call(ast.NumberTerm("100"), ast.NumberTerm("1"))),
				ast.NewExpr(ast.LessThan.Call(ast.NumberTerm("99"), ast.NumberTerm("1"))),
			),
		},
		want: `not {lt(100, 1); lt(99, 1)}`,
	},
	{
		name: "not, explicit body, multi-line",
		term: &ast.Not{
			ExplicitBody: true,
			Body: ast.NewBody(
				ast.NewExpr(ast.LessThan.Call(ast.NumberTerm("100"), ast.NumberTerm("1"))),
				ast.NewExpr(ast.LessThan.Call(ast.NumberTerm("99"), ast.NumberTerm("1"))),
			),
		},
		want: `not {lt(100, 1); lt(99, 1)}`,
	},
}

func TestASTValueTextAppendersAndStringLength(t *testing.T) {
	for _, tc := range valueAppenderTests {
		t.Run(tc.name, func(t *testing.T) {
			var buf []byte
			var err error

			buf, err = tc.term.(encoding.TextAppender).AppendText(buf)
			if err != nil {
				t.Fatalf("AppendText error: %v", err)
			}
			got := string(buf)
			if got != tc.want {
				t.Errorf("AppendText got %s, want %s", got, tc.want)
			}

			if expLen, gotLen := len(tc.want), tc.term.StringLength(); expLen != gotLen {
				t.Errorf("StringLength = %d; want %d", gotLen, expLen)
				t.Log(got)
			}
		})
	}
}

// Ensure no appender allocates when appending to a pre-sized buffer.
func BenchmarkNoASTTypeAllocatesOnAppendToBufferOfStringLength(b *testing.B) {
	for _, tc := range valueAppenderTests {
		b.Run(tc.name, func(b *testing.B) {
			buf := make([]byte, 0, tc.term.StringLength())
			for b.Loop() {
				_, err := tc.term.(encoding.TextAppender).AppendText(buf)
				if err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}
