/*
 * # Licensed to the LF AI & Data foundation 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 expr_test

import (
	"testing"

	"github.com/apache/arrow/go/v17/arrow/array"
	"github.com/apache/arrow/go/v17/arrow/memory"
	"github.com/stretchr/testify/suite"

	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
	"github.com/milvus-io/milvus/internal/util/function/chain"
	"github.com/milvus-io/milvus/internal/util/function/chain/expr"
	"github.com/milvus-io/milvus/internal/util/function/chain/types"
)

// =============================================================================
// Integration Test Suite
// =============================================================================

type DecayExprIntegrationTestSuite struct {
	suite.Suite
	pool *memory.CheckedAllocator
}

func (s *DecayExprIntegrationTestSuite) SetupTest() {
	s.pool = memory.NewCheckedAllocator(memory.NewGoAllocator())
}

func (s *DecayExprIntegrationTestSuite) TearDownTest() {
	s.pool.AssertSize(s.T(), 0)
}

func TestDecayExprIntegrationTestSuite(t *testing.T) {
	suite.Run(t, new(DecayExprIntegrationTestSuite))
}

// =============================================================================
// Helper Functions
// =============================================================================

func (s *DecayExprIntegrationTestSuite) createTestDataFrame(fieldType schemapb.DataType) *chain.DataFrame {
	var fieldData *schemapb.FieldData

	switch fieldType {
	case schemapb.DataType_Int64:
		fieldData = &schemapb.FieldData{
			Type:      schemapb.DataType_Int64,
			FieldName: "distance",
			FieldId:   100,
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_LongData{
						LongData: &schemapb.LongArray{
							Data: []int64{0, 50, 100, 150, 200, 0, 100, 200, 300},
						},
					},
				},
			},
		}
	case schemapb.DataType_Float:
		fieldData = &schemapb.FieldData{
			Type:      schemapb.DataType_Float,
			FieldName: "distance",
			FieldId:   100,
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_FloatData{
						FloatData: &schemapb.FloatArray{
							Data: []float32{0, 50, 100, 150, 200, 0, 100, 200, 300},
						},
					},
				},
			},
		}
	case schemapb.DataType_Double:
		fieldData = &schemapb.FieldData{
			Type:      schemapb.DataType_Double,
			FieldName: "distance",
			FieldId:   100,
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_DoubleData{
						DoubleData: &schemapb.DoubleArray{
							Data: []float64{0, 50, 100, 150, 200, 0, 100, 200, 300},
						},
					},
				},
			},
		}
	case schemapb.DataType_Int32:
		fieldData = &schemapb.FieldData{
			Type:      schemapb.DataType_Int32,
			FieldName: "distance",
			FieldId:   100,
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_IntData{
						IntData: &schemapb.IntArray{
							Data: []int32{0, 50, 100, 150, 200, 0, 100, 200, 300},
						},
					},
				},
			},
		}
	}

	resultData := &schemapb.SearchResultData{
		NumQueries: 2,
		TopK:       5,
		Topks:      []int64{5, 4},
		Scores:     []float32{1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0},
		Ids: &schemapb.IDs{
			IdField: &schemapb.IDs_IntId{
				IntId: &schemapb.LongArray{
					Data: []int64{1, 2, 3, 4, 5, 6, 7, 8, 9},
				},
			},
		},
		FieldsData: []*schemapb.FieldData{fieldData},
	}

	df, err := chain.FromSearchResultData(resultData, s.pool, []string{"distance"})
	s.Require().NoError(err)
	return df
}

// =============================================================================
// Integration Tests
// =============================================================================

func (s *DecayExprIntegrationTestSuite) TestIntegration_ChainWithDecay() {
	df := s.createTestDataFrame(schemapb.DataType_Int64)
	defer df.Release()

	decayExpr, err := expr.NewDecayExpr(expr.GaussFunction, 100, 50, 0, 0.5)
	s.Require().NoError(err)

	combineExpr, err := expr.NewNumCombineExpr(expr.ModeMultiply, nil)
	s.Require().NoError(err)

	// DecayExpr outputs pure decay factor into "_decay_score",
	// then NumCombineExpr multiplies $score with _decay_score.
	result, err := chain.NewFuncChainWithAllocator(s.pool).
		SetStage(types.StageL2Rerank).
		Map(decayExpr, []string{"distance"}, []string{"_decay_score"}).
		Map(combineExpr, []string{types.ScoreFieldName, "_decay_score"}, []string{types.ScoreFieldName}).
		Sort(types.ScoreFieldName, true, types.IDFieldName). // descending
		Limit(3).
		Execute(df)
	s.Require().NoError(err)
	defer result.Release()

	// Should have 3 results per chunk
	s.Equal([]int64{3, 3}, result.ChunkSizes())
}

func (s *DecayExprIntegrationTestSuite) TestIntegration_NullDecayFactorTreatedAsZero() {
	fieldData := &schemapb.FieldData{
		Type:      schemapb.DataType_Int64,
		FieldName: "distance",
		FieldId:   100,
		ValidData: []bool{true, false},
		Field: &schemapb.FieldData_Scalars{
			Scalars: &schemapb.ScalarField{
				Data: &schemapb.ScalarField_LongData{
					LongData: &schemapb.LongArray{
						Data: []int64{100, 0},
					},
				},
			},
		},
	}
	resultData := &schemapb.SearchResultData{
		NumQueries: 1,
		TopK:       2,
		Topks:      []int64{2},
		Scores:     []float32{1.0, 0.8},
		Ids: &schemapb.IDs{
			IdField: &schemapb.IDs_IntId{
				IntId: &schemapb.LongArray{
					Data: []int64{1, 2},
				},
			},
		},
		FieldsData: []*schemapb.FieldData{fieldData},
	}
	df, err := chain.FromSearchResultData(resultData, s.pool, []string{"distance"})
	s.Require().NoError(err)
	defer df.Release()

	decayExpr, err := expr.NewDecayExpr(expr.GaussFunction, 100, 50, 0, 0.5)
	s.Require().NoError(err)
	combineExpr, err := expr.NewNumCombineExpr(expr.ModeMultiply, nil, expr.WithNullPolicy(expr.NumCombineNullAsZero))
	s.Require().NoError(err)

	result, err := chain.NewFuncChainWithAllocator(s.pool).
		SetStage(types.StageL2Rerank).
		Map(decayExpr, []string{"distance"}, []string{"_decay_score"}).
		Map(combineExpr, []string{types.ScoreFieldName, "_decay_score"}, []string{types.ScoreFieldName}).
		Execute(df)
	s.Require().NoError(err)
	defer result.Release()

	scores := result.Column(types.ScoreFieldName).Chunk(0).(*array.Float32)
	s.False(scores.IsNull(0))
	s.InDelta(1.0, scores.Value(0), 0.001)
	s.False(scores.IsNull(1))
	s.InDelta(0.0, scores.Value(1), 0.001)
}

func (s *DecayExprIntegrationTestSuite) TestIntegration_ParseFromProto() {
	// DecayExpr only takes distance input, outputs decay factor into _decay_score.
	// NumCombineExpr then multiplies $score with _decay_score.
	fc, err := chain.ParseFuncChainProto(&schemapb.FunctionChain{
		Name:  "decay-test",
		Stage: schemapb.FunctionChainStage_FunctionChainStageL2Rerank,
		Ops: []*schemapb.FunctionChainOp{
			{
				Op:      types.OpTypeMap,
				Outputs: []string{"_decay_score"},
				Expr: &schemapb.FunctionChainExpr{
					Name: "decay",
					Args: []*schemapb.FunctionChainExprArg{
						columnArg("distance"),
					},
					Params: map[string]*schemapb.FunctionParamValue{
						types.DecayParamFunction: stringParam(types.DecayFuncGauss),
						types.DecayParamOrigin:   intParam(100),
						types.DecayParamScale:    intParam(50),
						types.DecayParamDecay:    doubleParam(0.5),
					},
				},
			},
			{
				Op:      types.OpTypeMap,
				Outputs: []string{types.ScoreFieldName},
				Expr: &schemapb.FunctionChainExpr{
					Name: "num_combine",
					Args: []*schemapb.FunctionChainExprArg{
						columnArg(types.ScoreFieldName),
						columnArg("_decay_score"),
					},
					Params: map[string]*schemapb.FunctionParamValue{
						types.NumCombineParamMode: stringParam(types.NumCombineModeMultiply),
					},
				},
			},
		},
	}, s.pool)
	s.Require().NoError(err)

	df := s.createTestDataFrame(schemapb.DataType_Int64)
	defer df.Release()

	result, err := fc.Execute(df)
	s.Require().NoError(err)
	defer result.Release()

	s.True(result.HasColumn(types.ScoreFieldName))
}

func (s *DecayExprIntegrationTestSuite) TestIntegration_NumCombine_ParseFromProto() {
	fc, err := chain.ParseFuncChainProto(&schemapb.FunctionChain{
		Name:  "num-combine-test",
		Stage: schemapb.FunctionChainStage_FunctionChainStageL2Rerank,
		Ops: []*schemapb.FunctionChainOp{
			{
				Op:      types.OpTypeMap,
				Outputs: []string{types.ScoreFieldName},
				Expr: &schemapb.FunctionChainExpr{
					Name: "num_combine",
					Args: []*schemapb.FunctionChainExprArg{
						columnArg(types.ScoreFieldName),
						columnArg("_func_score"),
					},
					Params: map[string]*schemapb.FunctionParamValue{
						types.NumCombineParamMode: stringParam(types.NumCombineModeMultiply),
					},
				},
			},
		},
	}, s.pool)
	s.Require().NoError(err)
	s.NotNil(fc)
}

func columnArg(name string) *schemapb.FunctionChainExprArg {
	return &schemapb.FunctionChainExprArg{Arg: &schemapb.FunctionChainExprArg_Column{Column: &schemapb.FunctionChainColumnArg{Name: name}}}
}

func intParam(value int64) *schemapb.FunctionParamValue {
	return &schemapb.FunctionParamValue{Value: &schemapb.FunctionParamValue_Int64Value{Int64Value: value}}
}

func doubleParam(value float64) *schemapb.FunctionParamValue {
	return &schemapb.FunctionParamValue{Value: &schemapb.FunctionParamValue_DoubleValue{DoubleValue: value}}
}

func stringParam(value string) *schemapb.FunctionParamValue {
	return &schemapb.FunctionParamValue{Value: &schemapb.FunctionParamValue_StringValue{StringValue: value}}
}