//                           _       _
// __      _____  __ ___   ___  __ _| |_ ___
// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
//  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
//   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
//
//  Copyright © 2016 - 2026 Weaviate B.V. All rights reserved.
//
//  CONTACT: hello@weaviate.io
//

package db

import (
	"context"
	"hash/crc32"
	"io"
	"slices"
	"testing"

	"github.com/go-openapi/strfmt"
	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	"github.com/weaviate/weaviate/entities/loadlimiter"

	"github.com/weaviate/weaviate/adapters/repos/db/inverted"
	"github.com/weaviate/weaviate/adapters/repos/db/queue"
	"github.com/weaviate/weaviate/adapters/repos/db/roaringset"
	resolver "github.com/weaviate/weaviate/adapters/repos/db/sharding"
	"github.com/weaviate/weaviate/entities/models"
	"github.com/weaviate/weaviate/entities/schema"
	"github.com/weaviate/weaviate/entities/storagestate"
	"github.com/weaviate/weaviate/entities/storobj"
	"github.com/weaviate/weaviate/entities/vectorindex/hnsw"
	"github.com/weaviate/weaviate/usecases/memwatch"
	"github.com/weaviate/weaviate/usecases/monitoring"
	schemaUC "github.com/weaviate/weaviate/usecases/schema"
	"github.com/weaviate/weaviate/usecases/sharding"
)

func TestUpdateIndexTenants(t *testing.T) {
	tests := []struct {
		name           string
		originalStatus string
		incomingStatus string
		expectedStatus storagestate.Status
		getClass       bool
	}{
		{
			name:           "when tenant is marked as COLD in incoming state while being HOT in original index",
			originalStatus: models.TenantActivityStatusHOT,
			incomingStatus: models.TenantActivityStatusCOLD,
			expectedStatus: storagestate.StatusShutdown,
		},
		{
			name:           "when tenant is marked as HOT in incoming state while being COLD in original index",
			originalStatus: models.TenantActivityStatusCOLD,
			incomingStatus: models.TenantActivityStatusHOT,
			expectedStatus: storagestate.StatusReady,
			getClass:       true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mockSchemaGetter := schemaUC.NewMockSchemaGetter(t)
			mockSchemaGetter.On("NodeName").Return("node1").Maybe()

			class := &models.Class{
				Class:               "TestClass",
				InvertedIndexConfig: &models.InvertedIndexConfig{},
				MultiTenancyConfig: &models.MultiTenancyConfig{
					Enabled: true,
				},
			}
			if tt.getClass {
				mockSchemaGetter.On("ReadOnlyClass", "TestClass").Return(class)
			}
			logger := logrus.New()
			scheduler := queue.NewScheduler(queue.SchedulerOptions{
				Logger:  logger,
				Workers: 1,
			})

			// Create original index state
			originalSS := &sharding.State{
				Physical: map[string]sharding.Physical{
					"shard1": {
						Name:           "shard1",
						BelongsToNodes: []string{"node1"},
						Status:         tt.originalStatus,
					},
				},
				PartitioningEnabled: true,
			}

			mockSchemaReader := schemaUC.NewMockSchemaReader(t)
			mockSchemaReader.EXPECT().Read(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(className string, retryIfClassNotFound bool, readFunc func(*models.Class, *sharding.State) error) error {
				return readFunc(class, originalSS)
			}).Maybe()
			shardResolver := resolver.NewShardResolver(class.Class, class.MultiTenancyConfig.Enabled, mockSchemaGetter)
			index, err := NewIndex(context.Background(), IndexConfig{
				ClassName:         schema.ClassName("TestClass"),
				RootPath:          t.TempDir(),
				ReplicationFactor: 1,
				ShardLoadLimiter:  loadlimiter.NewLoadLimiter(monitoring.NoopRegisterer, "dummy", 1),
			}, inverted.ConfigFromModel(class.InvertedIndexConfig),
				hnsw.NewDefaultUserConfig(), nil, nil, shardResolver, mockSchemaGetter, mockSchemaReader, nil, logger, nil, nil, nil, nil, nil, class, nil, scheduler, nil, nil, NewShardReindexerV3Noop(), roaringset.NewBitmapBufPoolNoop(), false)
			require.NoError(t, err)

			shard, err := NewShard(context.Background(), nil, "shard1", index, class, nil, scheduler, nil,
				NewShardReindexerV3Noop(), false, roaringset.NewBitmapBufPoolNoop())
			require.NoError(t, err)

			index.shards.Store("shard1", shard)

			migrator := &Migrator{
				db: &DB{
					schemaGetter: mockSchemaGetter,
				},
				nodeId: "node1",
			}

			// Create incoming state
			incomingSS := &sharding.State{
				Physical: map[string]sharding.Physical{
					"shard1": {
						Name:           "shard1",
						BelongsToNodes: []string{"node1"},
						Status:         tt.incomingStatus,
					},
				},
				PartitioningEnabled: true,
			}

			err = migrator.updateIndexTenants(context.Background(), index, incomingSS)
			require.NoError(t, err)

			mockSchemaGetter.AssertExpectations(t)

			// Verify the shard status
			require.Equal(t, tt.expectedStatus, shard.GetStatus())
		})
	}
}

func TestUpdateIndexShards(t *testing.T) {
	tests := []struct {
		name           string
		initialShards  []string
		newShards      []string
		expectedShards []string
		mustLoad       bool
		lazyLoading    bool
	}{
		{
			name:           "add new shard with lazy loading",
			initialShards:  []string{"shard1", "shard2"},
			newShards:      []string{"shard1", "shard2", "shard3"},
			expectedShards: []string{"shard1", "shard2", "shard3"},
			mustLoad:       false,
			lazyLoading:    false,
		},
		{
			name:           "remove shard with lazy loading",
			initialShards:  []string{"shard1", "shard2", "shard3"},
			newShards:      []string{"shard1", "shard3"},
			expectedShards: []string{"shard1", "shard3"},
			mustLoad:       false,
			lazyLoading:    false,
		},
		{
			name:           "keep existing shards with lazy loading",
			initialShards:  []string{"shard1", "shard3"},
			newShards:      []string{"shard1", "shard3"},
			expectedShards: []string{"shard1", "shard3"},
			mustLoad:       false,
			lazyLoading:    false,
		},
		{
			name:           "add new shard with immediate loading",
			initialShards:  []string{"shard1", "shard2"},
			newShards:      []string{"shard1", "shard2", "shard3"},
			expectedShards: []string{"shard1", "shard2", "shard3"},
			mustLoad:       true,
			lazyLoading:    false,
		},
		{
			name:           "remove shard with immediate loading",
			initialShards:  []string{"shard1", "shard2", "shard3"},
			newShards:      []string{"shard1", "shard3"},
			expectedShards: []string{"shard1", "shard3"},
			mustLoad:       true,
			lazyLoading:    false,
		},
		{
			name:           "keep existing shards with immediate loading",
			initialShards:  []string{"shard1", "shard3"},
			newShards:      []string{"shard1", "shard3"},
			expectedShards: []string{"shard1", "shard3"},
			mustLoad:       true,
			lazyLoading:    false,
		},
		{
			name:           "add new shard with lazy loading enabled",
			initialShards:  []string{"shard1", "shard2"},
			newShards:      []string{"shard1", "shard2", "shard3"},
			expectedShards: []string{"shard1", "shard2", "shard3"},
			mustLoad:       false,
			lazyLoading:    true,
		},
		{
			name:           "remove shard with lazy loading enabled",
			initialShards:  []string{"shard1", "shard2", "shard3"},
			newShards:      []string{"shard1", "shard3"},
			expectedShards: []string{"shard1", "shard3"},
			mustLoad:       false,
			lazyLoading:    true,
		},
		{
			name:           "keep existing shards with lazy loading enabled",
			initialShards:  []string{"shard1", "shard3"},
			newShards:      []string{"shard1", "shard3"},
			expectedShards: []string{"shard1", "shard3"},
			mustLoad:       false,
			lazyLoading:    true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := context.Background()
			logger := logrus.New()

			mockSchemaGetter := schemaUC.NewMockSchemaGetter(t)
			mockSchemaGetter.On("NodeName").Return("node1").Maybe()

			// Create a test class
			class := &models.Class{
				Class:               "TestClass",
				InvertedIndexConfig: &models.InvertedIndexConfig{},
				MultiTenancyConfig: &models.MultiTenancyConfig{
					Enabled: true,
				},
			}
			mockSchemaGetter.On("ReadOnlyClass", "TestClass").Return(class).Maybe()

			// Create initial sharding state
			initialPhysical := make(map[string]sharding.Physical)
			for _, shard := range tt.initialShards {
				initialPhysical[shard] = sharding.Physical{
					Name:           shard,
					BelongsToNodes: []string{"node1"},
				}
			}
			initialState := &sharding.State{
				Physical: initialPhysical,
			}
			initialState.SetLocalName("node1")
			scheduler := queue.NewScheduler(queue.SchedulerOptions{
				Logger:  logger,
				Workers: 1,
			})
			mockSchemaReader := schemaUC.NewMockSchemaReader(t)
			mockSchemaReader.EXPECT().Read(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(className string, retryIfClassNotFound bool, readFunc func(*models.Class, *sharding.State) error) error {
				return readFunc(class, initialState)
			}).Maybe()
			shardResolver := resolver.NewShardResolver(class.Class, class.MultiTenancyConfig.Enabled, mockSchemaGetter)
			// Create index with proper configuration
			index, err := NewIndex(ctx, IndexConfig{
				ClassName:             schema.ClassName("TestClass"),
				RootPath:              t.TempDir(),
				ReplicationFactor:     1,
				ShardLoadLimiter:      loadlimiter.NewLoadLimiter(monitoring.NoopRegisterer, "dummy", 1),
				DisableLazyLoadShards: !tt.lazyLoading, // Enable lazy loading when lazyLoading is true
			}, inverted.ConfigFromModel(class.InvertedIndexConfig),
				hnsw.NewDefaultUserConfig(), nil, nil, shardResolver, mockSchemaGetter, mockSchemaReader, nil, logger, nil, nil, nil, nil, nil, class, nil, scheduler, nil, memwatch.NewDummyMonitor(), NewShardReindexerV3Noop(), roaringset.NewBitmapBufPoolNoop(), false)
			require.NoError(t, err)

			// Initialize shards
			for _, shardName := range tt.initialShards {
				err := index.initLocalShardWithForcedLoading(ctx, class, shardName, tt.mustLoad, false)
				require.NoError(t, err)
			}

			migrator := &Migrator{
				db: &DB{
					schemaGetter: mockSchemaGetter,
				},
				nodeId: "node1",
			}

			// Create new sharding state
			newPhysical := make(map[string]sharding.Physical)
			for _, shard := range tt.newShards {
				newPhysical[shard] = sharding.Physical{
					Name:           shard,
					BelongsToNodes: []string{"node1"},
				}
			}
			newState := &sharding.State{
				Physical: newPhysical,
			}
			newState.SetLocalName("node1")

			// Update shards
			err = migrator.updateIndexShards(ctx, index, newState)
			require.NoError(t, err)

			// Verify expected shards exist and are of the correct type and status
			for _, expectedShard := range tt.expectedShards {
				shard := index.shards.Load(expectedShard)
				require.NotNil(t, shard, "shard %s should exist", expectedShard)

				_, isLazy := shard.(*LazyLoadShard)
				if tt.lazyLoading {
					// If lazyLoading is true, shard should be a LazyLoadShard
					require.True(t, isLazy, "shard %s should be a LazyLoadShard when lazyLoading=true", expectedShard)
					status := shard.GetStatus()
					require.True(t, status == storagestate.StatusLazyLoading, "shard %s should be in lazy loading state", expectedShard)
				} else {
					require.False(t, isLazy, "shard %s should be a regular Shard when lazyLoading=false", expectedShard)
					require.Equal(t, storagestate.StatusReady, shard.GetStatus(), "shard %s should be ready", expectedShard)
				}
			}

			// Verify removed shards are dropped
			for _, initialShard := range tt.initialShards {
				if !slices.Contains(tt.newShards, initialShard) {
					shard := index.shards.Load(initialShard)
					require.Nil(t, shard, "shard %s should be dropped", initialShard)
				}
			}

			mockSchemaGetter.AssertExpectations(t)
		})
	}
}

func TestListAndGetFilesWithIntegrityChecking(t *testing.T) {
	mockSchemaGetter := schemaUC.NewMockSchemaGetter(t)
	mockSchemaGetter.On("NodeName").Return("node1")

	class := &models.Class{
		Class:               "TestClass",
		InvertedIndexConfig: &models.InvertedIndexConfig{},
		MultiTenancyConfig: &models.MultiTenancyConfig{
			Enabled: true,
		},
	}
	mockSchemaGetter.On("ReadOnlyClass", "TestClass").Return(class).Maybe()

	logger := logrus.New()
	scheduler := queue.NewScheduler(queue.SchedulerOptions{
		Logger:  logger,
		Workers: 1,
	})

	// Create original index state
	originalSS := &sharding.State{
		Physical: map[string]sharding.Physical{
			"shard1": {
				Name:           "shard1",
				BelongsToNodes: []string{"node1"},
				Status:         models.TenantActivityStatusHOT,
			},
		},
		PartitioningEnabled: true,
	}

	mockSchemaReader := schemaUC.NewMockSchemaReader(t)
	mockSchemaReader.EXPECT().Read(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(className string, retryIfClassNotFound bool, readFunc func(*models.Class, *sharding.State) error) error {
		return readFunc(class, originalSS)
	}).Maybe()
	shardResolver := resolver.NewShardResolver(class.Class, class.MultiTenancyConfig.Enabled, mockSchemaGetter)
	index, err := NewIndex(context.Background(), IndexConfig{
		ClassName:         schema.ClassName("TestClass"),
		RootPath:          t.TempDir(),
		ReplicationFactor: 1,
		ShardLoadLimiter:  loadlimiter.NewLoadLimiter(monitoring.NoopRegisterer, "dummy", 1),
	}, inverted.ConfigFromModel(class.InvertedIndexConfig),
		hnsw.NewDefaultUserConfig(), nil, nil, shardResolver, mockSchemaGetter, mockSchemaReader, nil, logger, nil, nil, nil, nil, nil, class, nil, scheduler, nil, nil, NewShardReindexerV3Noop(), roaringset.NewBitmapBufPoolNoop(), false)
	require.NoError(t, err)

	shard, err := NewShard(context.Background(), nil, "shard1", index, class, nil, scheduler, nil,
		NewShardReindexerV3Noop(), false, roaringset.NewBitmapBufPoolNoop())
	require.NoError(t, err)

	index.shards.Store("shard1", shard)

	ctx := context.Background()

	err = index.IncomingPutObject(ctx, "shard1", &storobj.Object{
		MarshallerVersion: 1,
		DocID:             0,
		Object: models.Object{
			ID:    strfmt.UUID("40d3be3e-2ecc-49c8-b37c-d8983164848b"),
			Class: "TestClass",
		},
	}, 0)
	require.NoError(t, err)

	err = index.IncomingPauseFileActivity(ctx, "shard1")
	require.NoError(t, err)

	files, err := index.IncomingListFiles(ctx, "shard1")
	require.NoError(t, err)
	require.NotEmpty(t, files)

	for i, f := range files {
		md, err := index.IncomingGetFileMetadata(ctx, "shard1", f)
		require.NoError(t, err)

		// object insertion should not affect file copy process
		err = index.IncomingPutObject(ctx, "shard1", &storobj.Object{
			MarshallerVersion: 1,
			DocID:             uint64(i) + 1,
			Object: models.Object{
				ID:    strfmt.UUID("40d3be3e-2ecc-49c8-b37c-d8983164848b"),
				Class: "TestClass",
			},
		}, 0)
		require.NoError(t, err)

		r, err := index.IncomingGetFile(ctx, "shard1", f)
		require.NoError(t, err)

		h := crc32.NewIEEE()

		_, err = io.Copy(h, r)
		require.NoError(t, err)

		require.Equal(t, md.CRC32, h.Sum32())
	}

	err = index.IncomingResumeFileActivity(ctx, "shard1")
	require.NoError(t, err)
}
