// Licensed to 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. Apache Software Foundation (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 stream

import (
	"context"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"github.com/stretchr/testify/require"

	"github.com/apache/skywalking-banyandb/api/common"
	itest "github.com/apache/skywalking-banyandb/banyand/internal/test"
	"github.com/apache/skywalking-banyandb/banyand/protector"
	"github.com/apache/skywalking-banyandb/pkg/fs"
	"github.com/apache/skywalking-banyandb/pkg/logger"
	pbv1 "github.com/apache/skywalking-banyandb/pkg/pb/v1"
	"github.com/apache/skywalking-banyandb/pkg/query/model"
	"github.com/apache/skywalking-banyandb/pkg/test"
	"github.com/apache/skywalking-banyandb/pkg/timestamp"
)

func TestQueryResult_QuotaExceeded(t *testing.T) {
	type testCtx struct {
		name                string
		esList              []*elements
		sids                []common.SeriesID
		want                []blockMetadata
		minTimestamp        int64
		maxTimestamp        int64
		expectQuotaExceeded bool
	}

	tests := []testCtx{
		{
			name:                "TestQuotaNotExceeded_ExpectSuccess",
			esList:              []*elements{esTS1},
			sids:                []common.SeriesID{1, 2, 3},
			minTimestamp:        1,
			maxTimestamp:        1,
			expectQuotaExceeded: false,
			want: []blockMetadata{
				{seriesID: 1, count: 1, uncompressedSizeBytes: 859},
				{seriesID: 2, count: 1, uncompressedSizeBytes: 47},
				{seriesID: 3, count: 1, uncompressedSizeBytes: 16},
			},
		},
		{
			name:                "TestQuotaExceeded_ExpectError",
			esList:              []*elements{esTS1},
			sids:                []common.SeriesID{1, 2, 3},
			minTimestamp:        1,
			maxTimestamp:        1,
			expectQuotaExceeded: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Initialize a tstIter object.
			tmpPath, defFn := test.Space(require.New(t))
			fileSystem := fs.NewLocalFileSystem()
			defer defFn()
			tst, err := newTSTable(fileSystem, tmpPath, common.Position{},
				logger.GetLogger("test"), timestamp.TimeRange{}, option{flushTimeout: 0, mergePolicy: newDefaultMergePolicyForTesting(), protector: protector.Nop{}}, nil)
			require.NoError(t, err)
			for _, es := range tt.esList {
				tst.mustAddElements(es)
				time.Sleep(100 * time.Millisecond)
			}
			// wait until the introducer is done
			if len(tt.esList) > 0 {
				for {
					snp := tst.currentSnapshot()
					if snp == nil {
						time.Sleep(100 * time.Millisecond)
						continue
					}
					if snp.creator == snapshotCreatorMemPart {
						snp.decRef()
						time.Sleep(100 * time.Millisecond)
						continue
					}
					snp.decRef()
					tst.Close()
					break
				}
			}

			// reopen the table
			tst, err = newTSTable(fileSystem, tmpPath, common.Position{},
				logger.GetLogger("test"), timestamp.TimeRange{}, option{
					flushTimeout: defaultFlushTimeout, mergePolicy: newDefaultMergePolicyForTesting(),
					protector: protector.Nop{},
				}, nil)
			require.NoError(t, err)

			defer tst.Close()
			qr := &idxResult{
				pm:   &itest.MockMemoryProtector{ExpectQuotaExceeded: tt.expectQuotaExceeded},
				tabs: []*tsTable{tst},
			}
			qo := queryOptions{
				minTimestamp: tt.minTimestamp,
				maxTimestamp: tt.maxTimestamp,
				sortedSids:   tt.sids,
			}
			err = qr.scanParts(context.TODO(), qo)
			if tt.expectQuotaExceeded {
				require.Error(t, err)
				require.Contains(t, err.Error(), "quota exceeded", "expected quota to be exceeded but got: %v", err)
				return
			}
			require.NoError(t, err)
			var got []blockMetadata
			for _, data := range qr.data {
				got = append(got, data.bm)
			}
			if diff := cmp.Diff(got, tt.want,
				cmpopts.IgnoreFields(blockMetadata{}, "timestamps"),
				cmpopts.IgnoreFields(blockMetadata{}, "elementIDs"),
				cmpopts.IgnoreFields(blockMetadata{}, "tagFamilies"),
				cmp.AllowUnexported(blockMetadata{}),
			); diff != "" {
				t.Errorf("Unexpected blockMetadata (-got +want):\n%s", diff)
			}
		})
	}
}

func TestIdxResult_MergeByTagValue_Deduplication(t *testing.T) {
	tests := []struct {
		name             string
		data             []*blockCursor
		elementIDsSorted []uint64
		wantElementIDs   []uint64
		wantLen          int
	}{
		{
			name: "duplicate element IDs in block cursors should be deduplicated",
			data: []*blockCursor{
				{
					timestamps: []int64{10, 20},
					elementIDs: []uint64{100, 200},
					idx:        0,
					bm:         blockMetadata{seriesID: common.SeriesID(1)},
					tagProjection: []model.TagProjection{
						{Family: "family1", Names: []string{"tag1"}},
					},
					tagFamilies: []tagFamily{
						{
							name: "family1",
							tags: []tag{
								{name: "tag1", values: [][]byte{[]byte("value1"), []byte("value2")}},
							},
						},
					},
					schemaTagTypes: make(map[string]pbv1.ValueType),
				},
				{
					timestamps: []int64{15, 25},
					elementIDs: []uint64{100, 300}, // 100 is duplicate
					idx:        0,
					bm:         blockMetadata{seriesID: common.SeriesID(2)},
					tagProjection: []model.TagProjection{
						{Family: "family1", Names: []string{"tag1"}},
					},
					tagFamilies: []tagFamily{
						{
							name: "family1",
							tags: []tag{
								{name: "tag1", values: [][]byte{[]byte("value1_dup"), []byte("value3")}},
							},
						},
					},
					schemaTagTypes: make(map[string]pbv1.ValueType),
				},
			},
			elementIDsSorted: []uint64{100, 200, 300},
			wantElementIDs:   []uint64{100, 200, 300},
			wantLen:          3,
		},
		{
			name: "duplicate element IDs in elementIDsSorted should be deduplicated",
			data: []*blockCursor{
				{
					timestamps: []int64{10, 20, 30},
					elementIDs: []uint64{100, 200, 300},
					idx:        0,
					bm:         blockMetadata{seriesID: common.SeriesID(1)},
					tagProjection: []model.TagProjection{
						{Family: "family1", Names: []string{"tag1"}},
					},
					tagFamilies: []tagFamily{
						{
							name: "family1",
							tags: []tag{
								{name: "tag1", values: [][]byte{[]byte("value1"), []byte("value2"), []byte("value3")}},
							},
						},
					},
					schemaTagTypes: make(map[string]pbv1.ValueType),
				},
			},
			elementIDsSorted: []uint64{100, 200, 100, 300}, // 100 appears twice
			wantElementIDs:   []uint64{100, 200, 300},      // Should be deduplicated
			wantLen:          3,
		},
		{
			name: "duplicate element IDs in both data and elementIDsSorted",
			data: []*blockCursor{
				{
					timestamps: []int64{10, 20},
					elementIDs: []uint64{100, 200},
					idx:        0,
					bm:         blockMetadata{seriesID: common.SeriesID(1)},
					tagProjection: []model.TagProjection{
						{Family: "family1", Names: []string{"tag1"}},
					},
					tagFamilies: []tagFamily{
						{
							name: "family1",
							tags: []tag{
								{name: "tag1", values: [][]byte{[]byte("value1"), []byte("value2")}},
							},
						},
					},
					schemaTagTypes: make(map[string]pbv1.ValueType),
				},
				{
					timestamps: []int64{15},
					elementIDs: []uint64{100}, // duplicate in data
					idx:        0,
					bm:         blockMetadata{seriesID: common.SeriesID(2)},
					tagProjection: []model.TagProjection{
						{Family: "family1", Names: []string{"tag1"}},
					},
					tagFamilies: []tagFamily{
						{
							name: "family1",
							tags: []tag{
								{name: "tag1", values: [][]byte{[]byte("value1_dup")}},
							},
						},
					},
					schemaTagTypes: make(map[string]pbv1.ValueType),
				},
			},
			elementIDsSorted: []uint64{100, 200, 100}, // 100 appears twice in sorted list
			wantElementIDs:   []uint64{100, 200},      // Should be deduplicated
			wantLen:          2,
		},
		{
			name: "all duplicates should result in single element",
			data: []*blockCursor{
				{
					timestamps: []int64{10},
					elementIDs: []uint64{100},
					idx:        0,
					bm:         blockMetadata{seriesID: common.SeriesID(1)},
					tagProjection: []model.TagProjection{
						{Family: "family1", Names: []string{"tag1"}},
					},
					tagFamilies: []tagFamily{
						{
							name: "family1",
							tags: []tag{
								{name: "tag1", values: [][]byte{[]byte("value1")}},
							},
						},
					},
					schemaTagTypes: make(map[string]pbv1.ValueType),
				},
				{
					timestamps: []int64{20},
					elementIDs: []uint64{100}, // duplicate
					idx:        0,
					bm:         blockMetadata{seriesID: common.SeriesID(2)},
					tagProjection: []model.TagProjection{
						{Family: "family1", Names: []string{"tag1"}},
					},
					tagFamilies: []tagFamily{
						{
							name: "family1",
							tags: []tag{
								{name: "tag1", values: [][]byte{[]byte("value1_dup")}},
							},
						},
					},
					schemaTagTypes: make(map[string]pbv1.ValueType),
				},
			},
			elementIDsSorted: []uint64{100, 100, 100}, // all duplicates
			wantElementIDs:   []uint64{100},           // Should result in single element
			wantLen:          1,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			qr := &idxResult{
				data:             tt.data,
				elementIDsSorted: tt.elementIDsSorted,
			}
			result := qr.mergeByTagValue()

			require.Equal(t, tt.wantLen, result.Len(), "unexpected length")
			require.Equal(t, tt.wantElementIDs, result.ElementIDs, "unexpected element IDs - duplicates should be removed")
			// Verify no duplicate element IDs
			seen := make(map[uint64]bool)
			for _, id := range result.ElementIDs {
				require.False(t, seen[id], "duplicate element ID found: %d", id)
				seen[id] = true
			}
		})
	}
}
