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

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"math/rand"
	"os"
	"path/filepath"
	"strconv"
	"testing"

	"github.com/samber/lo"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"

	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
	"github.com/milvus-io/milvus/client/v3/entity"
	"github.com/milvus-io/milvus/client/v3/sbbf"
)

func TestNewBloomFilterBlob(t *testing.T) {
	t.Run("int64 members probe true with no false negatives", func(t *testing.T) {
		members := []int64{1, 2, 3, -7, 1 << 40}
		blob, err := NewBloomFilterBlob(members, 0.001)
		require.NoError(t, err)
		f, err := sbbf.Parse(blob)
		require.NoError(t, err)
		for _, v := range members {
			assert.True(t, f.TestInt64(v), "member %d must probe true", v)
		}
	})

	t.Run("string members probe true", func(t *testing.T) {
		members := []string{"a", "bb", "小明"}
		blob, err := NewBloomFilterBlob(members, 0.01)
		require.NoError(t, err)
		f, err := sbbf.Parse(blob)
		require.NoError(t, err)
		for _, s := range members {
			assert.True(t, f.TestString(s), "member %q must probe true", s)
		}
	})

	t.Run("wrong member type is rejected", func(t *testing.T) {
		_, err := NewBloomFilterBlob([]float64{1.0}, 0.001)
		require.Error(t, err)
	})

	t.Run("fpr out of range is rejected", func(t *testing.T) {
		_, err := NewBloomFilterBlob([]int64{1}, 0.5)
		require.Error(t, err)
	})

	t.Run("fpr out of range is rejected for varchar members too", func(t *testing.T) {
		_, err := NewBloomFilterBlob([]string{"a"}, 0.5)
		require.Error(t, err)
	})
}

// TestNewBloomFilterBlobMatchesCppFixture checks the public Go builder against
// a checked-in blob generated by Apache Arrow C++'s BlockSplitBloomFilter. The
// fixture has 100 deterministic pseudo-random int64 members; each must probe
// true from the Arrow blob and Go must reproduce its MBF1 bytes exactly.
func TestNewBloomFilterBlobMatchesCppFixture(t *testing.T) {
	type cppFixture struct {
		Generator string   `json:"generator"`
		Seed      uint64   `json:"seed"`
		FPR       float64  `json:"fpr"`
		IntValues []string `json:"int_values"`
		BlobHex   string   `json:"blob_hex"`
	}

	data, err := os.ReadFile(filepath.Join("..", "sbbf", "testdata", "cpp_generated_100_int64.json"))
	require.NoError(t, err, "C++ fixture missing; regenerate it with MILVUS_SBBF_REGEN_CPP_FIXTURE=1 and BloomFilterExprTest.CppGeneratedFixtureIsReproducible")

	var fixture cppFixture
	require.NoError(t, json.Unmarshal(data, &fixture))
	require.Equal(t, "apache_arrow_parquet_block_split_bloom_filter", fixture.Generator)
	require.Len(t, fixture.IntValues, 100)

	values := make([]int64, 0, len(fixture.IntValues))
	for _, value := range fixture.IntValues {
		v, err := strconv.ParseInt(value, 10, 64)
		require.NoError(t, err)
		values = append(values, v)
	}

	cppBlob, err := hex.DecodeString(fixture.BlobHex)
	require.NoError(t, err)
	cppFilter, err := sbbf.Parse(cppBlob)
	require.NoError(t, err)
	for _, value := range values {
		require.True(t, cppFilter.TestInt64(value), "C++ blob has a false negative for %d", value)
	}

	goBlob, err := NewBloomFilterBlob(values, fixture.FPR)
	require.NoError(t, err)
	require.Equal(t, cppBlob, []byte(goBlob), "Go output must be bit-identical to the C++ fixture")
}

// TestBloomFilterBlobTemplateMarshaling pins the wire contract: a BloomFilterBlob
// travels as a raw bytes template value (no base64), so the proxy receives the
// blob verbatim and embeds it without rebuilding.
func TestBloomFilterBlobTemplateMarshaling(t *testing.T) {
	blob, err := NewBloomFilterBlob([]int64{10, 20, 30}, 0.001)
	require.NoError(t, err)

	tv, err := any2TmplValue(blob)
	require.NoError(t, err)

	bv, ok := tv.GetVal().(*schemapb.TemplateValue_BytesVal)
	require.True(t, ok, "BloomFilterBlob must marshal to a raw bytes template value")
	assert.Equal(t, []byte(blob), bv.BytesVal, "template bytes must equal the raw blob bytes (no base64)")
}

// BloomFilterSuite drives the PUBLIC option API (NewBloomFilterBlob →
// WithTemplateParam → Search/Query) through a mock server and asserts the blob
// lands in the outgoing request as a raw bytes_val — not base64, not stringified,
// not an array — closing the end-to-end client wiring gap.
type BloomFilterSuite struct {
	MockSuiteBase

	schema *entity.Schema
}

func (s *BloomFilterSuite) SetupSuite() {
	s.MockSuiteBase.SetupSuite()
	s.schema = entity.NewSchema().
		WithField(entity.NewField().WithName("ID").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)).
		WithField(entity.NewField().WithName("Vector").WithDataType(entity.FieldTypeFloatVector).WithDim(128))
}

func (s *BloomFilterSuite) floatVectors() []entity.Vector {
	return []entity.Vector{
		entity.FloatVector(lo.RepeatBy(128, func(_ int) float32 {
			return rand.Float32()
		})),
	}
}

// assertBloomBytesVal asserts the "bf" template value is a raw bytes_val equal to
// the blob (and specifically a *schemapb.TemplateValue_BytesVal, never a string or
// array variant).
func (s *BloomFilterSuite) assertBloomBytesVal(tmpl map[string]*schemapb.TemplateValue, blob BloomFilterBlob) {
	tv, ok := tmpl["bf"]
	s.Require().True(ok, "outgoing request must carry the 'bf' template value")
	bv, ok := tv.GetVal().(*schemapb.TemplateValue_BytesVal)
	s.Require().Truef(ok, "'bf' must be a *schemapb.TemplateValue_BytesVal, got %T", tv.GetVal())
	s.Equal([]byte(blob), bv.BytesVal, "template bytes must equal the raw blob bytes (no base64/stringify)")
}

func (s *BloomFilterSuite) TestSearchBloomMatchInt64() {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	collectionName := fmt.Sprintf("coll_%s", s.randString(6))
	s.setupCache(collectionName, s.schema)

	blob, err := NewBloomFilterBlob([]int64{1, 2, 3, 42, 1 << 40}, sbbf.DefaultFPR)
	s.Require().NoError(err)

	s.mock.EXPECT().Search(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, sr *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
		s.Equal(collectionName, sr.GetCollectionName())
		s.Contains(sr.GetDsl(), "bloom_match(myfield, {bf})")
		s.assertBloomBytesVal(sr.GetExprTemplateValues(), blob)

		return &milvuspb.SearchResults{
			Status: &commonpb.Status{},
			Results: &schemapb.SearchResultData{
				NumQueries: 1,
				TopK:       1,
				FieldsData: []*schemapb.FieldData{
					s.getInt64FieldData("ID", []int64{1}),
				},
				Ids:    &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}},
				Scores: []float32{0.1},
				Topks:  []int64{1},
			},
		}, nil
	}).Once()

	_, err = s.client.Search(ctx, NewSearchOption(collectionName, 10, s.floatVectors()).
		WithFilter("bloom_match(myfield, {bf})").
		WithTemplateParam("bf", blob))
	s.NoError(err)
}

func (s *BloomFilterSuite) TestSearchBloomMatchVarchar() {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	collectionName := fmt.Sprintf("coll_%s", s.randString(6))
	s.setupCache(collectionName, s.schema)

	blob, err := NewBloomFilterBlob([]string{"a", "bb", "小明"}, sbbf.DefaultFPR)
	s.Require().NoError(err)

	s.mock.EXPECT().Search(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, sr *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
		s.Equal(collectionName, sr.GetCollectionName())
		s.Contains(sr.GetDsl(), "bloom_match(myfield, {bf})")
		s.assertBloomBytesVal(sr.GetExprTemplateValues(), blob)

		return &milvuspb.SearchResults{
			Status: &commonpb.Status{},
			Results: &schemapb.SearchResultData{
				NumQueries: 1,
				TopK:       1,
				FieldsData: []*schemapb.FieldData{
					s.getInt64FieldData("ID", []int64{1}),
				},
				Ids:    &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}},
				Scores: []float32{0.1},
				Topks:  []int64{1},
			},
		}, nil
	}).Once()

	_, err = s.client.Search(ctx, NewSearchOption(collectionName, 10, s.floatVectors()).
		WithFilter("bloom_match(myfield, {bf})").
		WithTemplateParam("bf", blob))
	s.NoError(err)
}

func (s *BloomFilterSuite) TestQueryBloomMatchInt64() {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	collectionName := fmt.Sprintf("coll_%s", s.randString(6))
	s.setupCache(collectionName, s.schema)

	blob, err := NewBloomFilterBlob([]int64{1, 2, 3, 42, 1 << 40}, sbbf.DefaultFPR)
	s.Require().NoError(err)

	s.mock.EXPECT().Query(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, qr *milvuspb.QueryRequest) (*milvuspb.QueryResults, error) {
		s.Equal(collectionName, qr.GetCollectionName())
		s.Contains(qr.GetExpr(), "bloom_match(myfield, {bf})")
		s.assertBloomBytesVal(qr.GetExprTemplateValues(), blob)

		return &milvuspb.QueryResults{
			Status: &commonpb.Status{},
			FieldsData: []*schemapb.FieldData{
				s.getInt64FieldData("ID", []int64{1}),
			},
		}, nil
	}).Once()

	_, err = s.client.Query(ctx, NewQueryOption(collectionName).
		WithFilter("bloom_match(myfield, {bf})").
		WithTemplateParam("bf", blob).
		WithOutputFields("ID"))
	s.NoError(err)
}

func TestBloomFilter(t *testing.T) {
	suite.Run(t, new(BloomFilterSuite))
}
