// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package array_test

import (
	"fmt"
	"reflect"
	"testing"

	"github.com/apache/arrow/go/arrow"
	"github.com/apache/arrow/go/arrow/array"
	"github.com/apache/arrow/go/arrow/memory"
)

func TestRecord(t *testing.T) {
	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
	defer mem.AssertSize(t, 0)

	schema := arrow.NewSchema(
		[]arrow.Field{
			arrow.Field{Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32},
			arrow.Field{Name: "f2-f64", Type: arrow.PrimitiveTypes.Float64},
		},
		nil,
	)
	col1 := func() array.Interface {
		ib := array.NewInt32Builder(mem)
		defer ib.Release()

		ib.AppendValues([]int32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, nil)
		return ib.NewInt32Array()
	}()
	defer col1.Release()

	col2 := func() array.Interface {
		b := array.NewFloat64Builder(mem)
		defer b.Release()

		b.AppendValues([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, nil)
		return b.NewFloat64Array()
	}()
	defer col2.Release()

	cols := []array.Interface{col1, col2}
	rec := array.NewRecord(schema, cols, -1)
	defer rec.Release()

	rec.Retain()
	rec.Release()

	if got, want := rec.Schema(), schema; !got.Equal(want) {
		t.Fatalf("invalid schema: got=%#v, want=%#v", got, want)
	}

	if got, want := rec.NumRows(), int64(10); got != want {
		t.Fatalf("invalid number of rows: got=%d, want=%d", got, want)
	}
	if got, want := rec.NumCols(), int64(2); got != want {
		t.Fatalf("invalid number of columns: got=%d, want=%d", got, want)
	}
	if got, want := rec.Columns()[0], cols[0]; got != want {
		t.Fatalf("invalid column: got=%q, want=%q", got, want)
	}
	if got, want := rec.Column(0), cols[0]; got != want {
		t.Fatalf("invalid column: got=%q, want=%q", got, want)
	}
	if got, want := rec.ColumnName(0), schema.Field(0).Name; got != want {
		t.Fatalf("invalid column name: got=%q, want=%q", got, want)
	}

	for _, tc := range []struct {
		i, j int64
		err  error
	}{
		{i: 0, j: 10, err: nil},
		{i: 1, j: 10, err: nil},
		{i: 1, j: 9, err: nil},
		{i: 0, j: 0, err: nil},
		{i: 1, j: 1, err: nil},
		{i: 10, j: 10, err: nil},
		{i: 1, j: 0, err: fmt.Errorf("arrow/array: index out of range")},
		{i: 1, j: 11, err: fmt.Errorf("arrow/array: index out of range")},
	} {
		t.Run(fmt.Sprintf("slice-%02d-%02d", tc.i, tc.j), func(t *testing.T) {
			if tc.err != nil {
				defer func() {
					e := recover()
					if e == nil {
						t.Fatalf("expected an error %q", tc.err)
					}
					switch err := e.(type) {
					case string:
						if err != tc.err.Error() {
							t.Fatalf("invalid panic message. got=%q, want=%q", err, tc.err)
						}
					case error:
						if err.Error() != tc.err.Error() {
							t.Fatalf("invalid panic message. got=%q, want=%q", err, tc.err)
						}
					default:
						t.Fatalf("invalid type for panic message: %T (err=%v)", err, err)
					}
				}()
			}
			sub := rec.NewSlice(tc.i, tc.j)
			defer sub.Release()

			if got, want := sub.NumRows(), tc.j-tc.i; got != want {
				t.Fatalf("invalid rec-slice number of rows: got=%d, want=%d", got, want)
			}
		})
	}

	for _, tc := range []struct {
		schema *arrow.Schema
		cols   []array.Interface
		rows   int64
		err    error
	}{
		{
			schema: schema,
			cols:   nil,
			rows:   -1,
			err:    fmt.Errorf("arrow/array: number of columns/fields mismatch"),
		},
		{
			schema: schema,
			cols:   cols[:1],
			rows:   0,
			err:    fmt.Errorf("arrow/array: number of columns/fields mismatch"),
		},
		{
			schema: arrow.NewSchema(
				[]arrow.Field{
					arrow.Field{Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32},
				},
				nil,
			),
			cols: cols,
			rows: 0,
			err:  fmt.Errorf("arrow/array: number of columns/fields mismatch"),
		},
		{
			schema: arrow.NewSchema(
				[]arrow.Field{
					arrow.Field{Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32},
					arrow.Field{Name: "f2-f64", Type: arrow.PrimitiveTypes.Int32},
				},
				nil,
			),
			cols: cols,
			rows: 0,
			err:  fmt.Errorf(`arrow/array: column "f2-f64" type mismatch: got=float64, want=int32`),
		},
		{
			schema: schema,
			cols:   cols,
			rows:   11,
			err:    fmt.Errorf(`arrow/array: mismatch number of rows in column "f1-i32": got=10, want=11`),
		},
		{
			schema: schema,
			cols:   cols,
			rows:   10,
			err:    nil,
		},
		{
			schema: schema,
			cols:   cols,
			rows:   3,
			err:    nil,
		},
		{
			schema: schema,
			cols:   cols,
			rows:   0,
			err:    nil,
		},
	} {
		t.Run("", func(t *testing.T) {
			if tc.err != nil {
				defer func() {
					e := recover()
					if e == nil {
						t.Fatalf("expected an error %q", tc.err)
					}
					switch err := e.(type) {
					case string:
						if err != tc.err.Error() {
							t.Fatalf("invalid panic message. got=%q, want=%q", err, tc.err)
						}
					case error:
						if err.Error() != tc.err.Error() {
							t.Fatalf("invalid panic message. got=%q, want=%q", err, tc.err)
						}
					default:
						t.Fatalf("invalid type for panic message: %T (err=%v)", err, err)
					}
				}()
			}
			rec := array.NewRecord(tc.schema, tc.cols, tc.rows)
			defer rec.Release()
			if got, want := rec.NumRows(), tc.rows; got != want {
				t.Fatalf("invalid number of rows: got=%d, want=%d", got, want)
			}
		})
	}
}

func TestRecordReader(t *testing.T) {
	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
	defer mem.AssertSize(t, 0)

	schema := arrow.NewSchema(
		[]arrow.Field{
			arrow.Field{Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32},
			arrow.Field{Name: "f2-f64", Type: arrow.PrimitiveTypes.Float64},
		},
		nil,
	)
	rec1 := func() array.Record {
		col1 := func() array.Interface {
			ib := array.NewInt32Builder(mem)
			defer ib.Release()

			ib.AppendValues([]int32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, nil)
			return ib.NewInt32Array()
		}()
		defer col1.Release()

		col2 := func() array.Interface {
			b := array.NewFloat64Builder(mem)
			defer b.Release()

			b.AppendValues([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, nil)
			return b.NewFloat64Array()
		}()
		defer col2.Release()

		cols := []array.Interface{col1, col2}
		return array.NewRecord(schema, cols, -1)
	}()
	defer rec1.Release()

	rec2 := func() array.Record {
		col1 := func() array.Interface {
			ib := array.NewInt32Builder(mem)
			defer ib.Release()

			ib.AppendValues([]int32{11, 12, 13, 14, 15, 16, 17, 18, 19, 20}, nil)
			return ib.NewInt32Array()
		}()
		defer col1.Release()

		col2 := func() array.Interface {
			b := array.NewFloat64Builder(mem)
			defer b.Release()

			b.AppendValues([]float64{11, 12, 13, 14, 15, 16, 17, 18, 19, 20}, nil)
			return b.NewFloat64Array()
		}()
		defer col2.Release()

		cols := []array.Interface{col1, col2}
		return array.NewRecord(schema, cols, -1)
	}()
	defer rec2.Release()

	recs := []array.Record{rec1, rec2}
	itr, err := array.NewRecordReader(schema, recs)
	if err != nil {
		t.Fatal(err)
	}
	defer itr.Release()

	itr.Retain()
	itr.Release()

	if got, want := itr.Schema(), schema; !got.Equal(want) {
		t.Fatalf("invalid schema. got=%#v, want=%#v", got, want)
	}

	n := 0
	for itr.Next() {
		n++
		if got, want := itr.Record(), recs[n-1]; !reflect.DeepEqual(got, want) {
			t.Fatalf("itr[%d], invalid record. got=%#v, want=%#v", n-1, got, want)
		}
	}

	if n != len(recs) {
		t.Fatalf("invalid number of iterations. got=%d, want=%d", n, len(recs))
	}

	for _, tc := range []struct {
		name   string
		schema *arrow.Schema
		err    error
	}{
		{
			name: "mismatch-name",
			schema: arrow.NewSchema(
				[]arrow.Field{
					arrow.Field{Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32},
					arrow.Field{Name: "f2-XXX", Type: arrow.PrimitiveTypes.Float64},
				},
				nil,
			),
			err: fmt.Errorf("arrow/array: mismatch schema"),
		},
		{
			name: "mismatch-type",
			schema: arrow.NewSchema(
				[]arrow.Field{
					arrow.Field{Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32},
					arrow.Field{Name: "f2-f64", Type: arrow.PrimitiveTypes.Int64},
				},
				nil,
			),
			err: fmt.Errorf("arrow/array: mismatch schema"),
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			itr, err := array.NewRecordReader(tc.schema, recs)
			if itr != nil {
				itr.Release()
			}
			if err == nil {
				t.Fatalf("expected an error: %v", tc.err)
			}
			if !reflect.DeepEqual(tc.err, err) {
				t.Fatalf("invalid error: got=%v, want=%v", err, tc.err)
			}
		})
	}
}

func TestRecordBuilder(t *testing.T) {
	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
	defer mem.AssertSize(t, 0)

	schema := arrow.NewSchema(
		[]arrow.Field{
			arrow.Field{Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32},
			arrow.Field{Name: "f2-f64", Type: arrow.PrimitiveTypes.Float64},
		},
		nil,
	)

	b := array.NewRecordBuilder(mem, schema)
	defer b.Release()

	b.Retain()
	b.Release()

	b.Field(0).(*array.Int32Builder).AppendValues([]int32{1, 2, 3, 4, 5, 6}, nil)
	b.Field(0).(*array.Int32Builder).AppendValues([]int32{7, 8, 9, 10}, nil)
	b.Field(1).(*array.Float64Builder).AppendValues([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, nil)

	rec := b.NewRecord()
	defer rec.Release()

	if got, want := rec.Schema(), schema; !got.Equal(want) {
		t.Fatalf("invalid schema: got=%#v, want=%#v", got, want)
	}

	if got, want := rec.NumRows(), int64(10); got != want {
		t.Fatalf("invalid number of rows: got=%d, want=%d", got, want)
	}
	if got, want := rec.NumCols(), int64(2); got != want {
		t.Fatalf("invalid number of columns: got=%d, want=%d", got, want)
	}
	if got, want := rec.ColumnName(0), schema.Field(0).Name; got != want {
		t.Fatalf("invalid column name: got=%q, want=%q", got, want)
	}
}
