// 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 typeutil

import (
	"log"
	"testing"
	"unsafe"

	"github.com/stretchr/testify/assert"

	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
	"github.com/milvus-io/milvus/pkg/v3/common"
)

func TestUint64(t *testing.T) {
	var i int64 = -1
	u := uint64(i)
	t.Log(i)
	t.Log(u)
}

func TestHash32_Uint64(t *testing.T) {
	var u uint64 = 0x12
	h, err := Hash32Uint64(u)
	assert.NoError(t, err)
	t.Log(h)

	h1, err := Hash32Int64(int64(u))
	assert.NoError(t, err)
	t.Log(h1)
	assert.Equal(t, h, h1)

	/* #nosec G103 */
	b := make([]byte, unsafe.Sizeof(u))
	b[0] = 0x12
	h2, err := Hash32Bytes(b)
	assert.NoError(t, err)

	t.Log(h2)
	assert.Equal(t, h, h2)
}

func TestHash32_String(t *testing.T) {
	u := "ok"
	h, err := Hash32String(u)
	assert.NoError(t, err)

	t.Log(h)
	log.Println(h)

	b := []byte(u)
	h2, err := Hash32Bytes(b)
	assert.NoError(t, err)
	log.Println(h2)

	assert.Equal(t, uint32(h), h2)
}

func TestHash32_Uint64_Value(t *testing.T) {
	h, _ := Hash32Int64(int64(0))
	assert.Equal(t, uint32(1669671676), h)

	h, _ = Hash32Int64(int64(1))
	assert.Equal(t, uint32(1392991556), h)

	h, _ = Hash32Int64(int64(10))
	assert.Equal(t, uint32(1857498428), h)

	h, _ = Hash32Int64(int64(100))
	assert.Equal(t, uint32(1177227376), h)

	h, _ = Hash32Int64(int64(1000))
	assert.Equal(t, uint32(679959820), h)

	h, _ = Hash32Int64(int64(10000))
	assert.Equal(t, uint32(1785225994), h)

	h, _ = Hash32Int64(int64(100000))
	assert.Equal(t, uint32(258819970), h)

	h, _ = Hash32Int64(int64(1000000))
	assert.Equal(t, uint32(1669527334), h)

	h, _ = Hash32Int64(int64(10000000))
	assert.Equal(t, uint32(112035536), h)

	h, _ = Hash32Int64(int64(100000000))
	assert.Equal(t, uint32(107125476), h)

	h, _ = Hash32Int64(int64(1000000000))
	assert.Equal(t, uint32(913068403), h)

	h, _ = Hash32Int64(int64(10000000000))
	assert.Equal(t, uint32(1151135), h)

	h, _ = Hash32Int64(int64(100000000000))
	assert.Equal(t, uint32(2136277055), h)

	h, _ = Hash32Int64(int64(1000000000000))
	assert.Equal(t, uint32(636570700), h)

	h, _ = Hash32Int64(int64(10000000000000))
	assert.Equal(t, uint32(469485073), h)

	h, _ = Hash32Int64(int64(100000000000000))
	assert.Equal(t, uint32(745519692), h)

	h, _ = Hash32Int64(int64(1000000000000000))
	assert.Equal(t, uint32(1724407835), h)

	h, _ = Hash32Int64(int64(10000000000000000))
	assert.Equal(t, uint32(2022854314), h)

	h, _ = Hash32Int64(int64(100000000000000000))
	assert.Equal(t, uint32(1713940411), h)

	h, _ = Hash32Int64(int64(1000000000000000000))
	assert.Equal(t, uint32(2117120843), h)
}

func TestHashString2Uint32(t *testing.T) {
	u := "abcdef"
	assert.Equal(t, uint32(1267612143), HashString2Uint32(u))

	longStr := "Milvus was created in 2019 with a singular goal: store, index, and manage massive embedding vectors generated by deep neural networks and other machine learning (ML) models.\n\nAs a database specifically designed to handle queries over input vectors, it is capable of indexing vectors on a trillion scale. Unlike existing relational databases which mainly deal with structured data following a pre-defined pattern, Milvus is designed from the bottom-up to handle embedding vectors converted from unstructured data.\n\nAs the Internet grew and evolved, unstructured data became more and more common, including emails, papers, IoT sensor data, Facebook photos, protein structures, and much more. In order for computers to understand and process unstructured data, these are converted into vectors using embedding techniques. Milvus stores and indexes these vectors. Milvus is able to analyze the correlation between two vectors by calculating their similarity distance. If the two embedding vectors are very similar, it means that the original data sources are similar as well."
	hash := HashString2Uint32(longStr)
	assert.Equal(t, uint32(2922592797), hash)
}

func TestHashPK2Channels(t *testing.T) {
	channels := []string{"test1", "test2"}
	int64IDs := &schemapb.IDs{
		IdField: &schemapb.IDs_IntId{
			IntId: &schemapb.LongArray{
				Data: []int64{100, 102, 102, 103, 104},
			},
		},
	}
	ret, err := HashPK2Channels(int64IDs, channels)
	assert.NoError(t, err)
	assert.Equal(t, 5, len(ret))
	// same pk hash to same channel
	assert.Equal(t, ret[1], ret[2])

	stringIDs := &schemapb.IDs{
		IdField: &schemapb.IDs_StrId{
			StrId: &schemapb.StringArray{
				Data: []string{"ab", "bc", "bc", "abd", "milvus"},
			},
		},
	}
	ret, err = HashPK2Channels(stringIDs, channels)
	assert.NoError(t, err)
	assert.Equal(t, 5, len(ret))
	assert.Equal(t, ret[1], ret[2])
}

func TestHashPK2ChannelsMatchesLegacyModulo(t *testing.T) {
	channels := []string{"channel-0", "channel-1", "channel-2", "channel-3"}

	t.Run("int64 primary keys", func(t *testing.T) {
		keys := []int64{0, 1, 10, 100, 1000, -1}
		ids := &schemapb.IDs{
			IdField: &schemapb.IDs_IntId{
				IntId: &schemapb.LongArray{Data: keys},
			},
		}
		expected := make([]uint32, 0, len(keys))
		for _, key := range keys {
			hash, err := Hash32Int64(key)
			assert.NoError(t, err)
			expected = append(expected, hash%uint32(len(channels)))
		}

		got, err := HashPK2Channels(ids, channels)

		assert.NoError(t, err)
		assert.Equal(t, expected, got)
	})

	t.Run("varchar primary keys", func(t *testing.T) {
		keys := []string{"ab", "bc", "bc", "abd", "milvus"}
		ids := &schemapb.IDs{
			IdField: &schemapb.IDs_StrId{
				StrId: &schemapb.StringArray{Data: keys},
			},
		}
		expected := make([]uint32, 0, len(keys))
		for _, key := range keys {
			expected = append(expected, HashString2Uint32(key)%uint32(len(channels)))
		}

		got, err := HashPK2Channels(ids, channels)

		assert.NoError(t, err)
		assert.Equal(t, expected, got)
	})
}

func TestHashPK2ChannelsReturnsRoutingError(t *testing.T) {
	ids := &schemapb.IDs{
		IdField: &schemapb.IDs_IntId{
			IntId: &schemapb.LongArray{Data: []int64{1}},
		},
	}

	got, err := HashPK2Channels(ids, nil)

	assert.ErrorIs(t, err, common.ErrRoutingTableNoValues)
	assert.Nil(t, got)
}

func TestLocateRoutingIndexesMatchesLegacyModulo(t *testing.T) {
	t.Run("int64 keys", func(t *testing.T) {
		keys := []int64{0, 1, 10, 100, 1000, -1}
		const targetCount = 5
		expected := make([]uint32, 0, len(keys))
		for _, key := range keys {
			hash, err := Hash32Int64(key)
			assert.NoError(t, err)
			expected = append(expected, hash%targetCount)
		}

		got, err := locateRoutingIndexes(keys, targetCount, int64RoutingHasher{})

		assert.NoError(t, err)
		assert.Equal(t, expected, got)
	})

	t.Run("varchar keys", func(t *testing.T) {
		keys := []string{"ab", "bc", "bc", "abd", "milvus"}
		const targetCount = 4
		expected := make([]uint32, 0, len(keys))
		for _, key := range keys {
			expected = append(expected, HashString2Uint32(key)%targetCount)
		}

		got, err := locateRoutingIndexes(keys, targetCount, stringRoutingHasher{})

		assert.NoError(t, err)
		assert.Equal(t, expected, got)
	})
}

func TestHashKey2PartitionsMatchesLegacyModulo(t *testing.T) {
	partitions := []string{"p_0", "p_1", "p_2"}

	t.Run("int64 keys", func(t *testing.T) {
		keys := []int64{0, 1, 10, 100, 1000, -1}
		fieldData := &schemapb.FieldData{
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_LongData{
						LongData: &schemapb.LongArray{Data: keys},
					},
				},
			},
		}
		expected := make([]uint32, 0, len(keys))
		for _, key := range keys {
			hash, err := Hash32Int64(key)
			assert.NoError(t, err)
			expected = append(expected, hash%uint32(len(partitions)))
		}

		got, err := HashKey2Partitions(fieldData, partitions)

		assert.NoError(t, err)
		assert.Equal(t, expected, got)
	})

	t.Run("varchar keys", func(t *testing.T) {
		keys := []string{"ab", "bc", "bc", "abd", "milvus"}
		fieldData := &schemapb.FieldData{
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_StringData{
						StringData: &schemapb.StringArray{Data: keys},
					},
				},
			},
		}
		expected := make([]uint32, 0, len(keys))
		for _, key := range keys {
			expected = append(expected, HashString2Uint32(key)%uint32(len(partitions)))
		}

		got, err := HashKey2Partitions(fieldData, partitions)

		assert.NoError(t, err)
		assert.Equal(t, expected, got)
	})
}

func TestRearrangePartitionsForPartitionKey(t *testing.T) {
	// invalid partition name
	partitions := map[string]int64{
		"invalid": 1,
	}

	partitionNames, partitionIDs, err := RearrangePartitionsForPartitionKey(partitions)
	assert.Error(t, err)
	assert.Nil(t, partitionNames)
	assert.Nil(t, partitionIDs)

	// invalid partition index
	partitions = map[string]int64{
		"invalid_a": 1,
	}

	partitionNames, partitionIDs, err = RearrangePartitionsForPartitionKey(partitions)
	assert.Error(t, err)
	assert.Nil(t, partitionNames)
	assert.Nil(t, partitionIDs)

	partitions = map[string]int64{
		"invalid_1": 1,
	}

	partitionNames, partitionIDs, err = RearrangePartitionsForPartitionKey(partitions)
	assert.Error(t, err)
	assert.Nil(t, partitionNames)
	assert.Nil(t, partitionIDs)

	// success cases
	validateFunc := func(partitions map[string]int64) {
		partitionNames, partitionIDs, err = RearrangePartitionsForPartitionKey(partitions)
		assert.NoError(t, err)
		assert.Equal(t, len(partitions), len(partitionNames))
		assert.Equal(t, len(partitions), len(partitionIDs))

		for i := 0; i < len(partitions); i++ {
			assert.Contains(t, partitions, partitionNames[i])
			assert.Equal(t, partitions[partitionNames[i]], partitionIDs[i])

			if i > 0 {
				assert.Greater(t, partitionIDs[i], partitionIDs[i-1])
			}
		}
	}

	validateFunc(map[string]int64{
		"p_0": 1,
		"p_1": 2,
		"p_2": 3,
	})

	validateFunc(map[string]int64{
		"p_2": 3,
		"p_1": 2,
		"p_0": 1,
	})

	validateFunc(map[string]int64{
		"p_1": 2,
		"p_2": 3,
		"p_0": 1,
	})
}
