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

package db

import (
	"context"
	"encoding/binary"
	"time"

	"github.com/pkg/errors"
	"github.com/weaviate/weaviate/adapters/repos/db/helpers"
	"github.com/weaviate/weaviate/adapters/repos/db/vector/common"
	"github.com/weaviate/weaviate/adapters/repos/db/vector/hnsw"
	"github.com/weaviate/weaviate/adapters/repos/db/vector/hnsw/visited"
	"github.com/weaviate/weaviate/entities/storobj"
)

// ConvertQueue converts a legacy in-memory queue to an on-disk queue.
// It detects if the queue has a checkpoint then it enqueues all the
// remaining vectors to the on-disk queue, then deletes the checkpoint.
func (s *Shard) ConvertQueue(targetVector string) error {
	if !asyncEnabled() {
		return nil
	}

	// load non-indexed vectors and add them to the queue
	checkpoint, exists, err := s.indexCheckpoints.Get(s.ID(), targetVector)
	if err != nil {
		return errors.Wrap(err, "get last indexed id")
	}
	if !exists {
		return nil
	}

	err = s.FillQueue(targetVector, checkpoint)
	if err != nil {
		return errors.WithStack(err)
	}

	// we can now safely remove the checkpoint
	err = s.indexCheckpoints.Delete(s.ID(), targetVector)
	if err != nil {
		return errors.Wrap(err, "delete checkpoint")
	}

	return nil
}

// FillQueue is a helper function that enqueues all vectors from the
// LSM store to the on-disk queue.
func (s *Shard) FillQueue(targetVector string, from uint64) error {
	if !asyncEnabled() {
		return nil
	}

	start := time.Now()

	var counter int

	vectorIndex, err := s.getVectorIndex(targetVector)
	if err != nil {
		s.index.logger.Warn("preload queue: vector index not found")
		// shard was never initialized, possibly because of a failed shard
		// initialization. No op.
		return nil
	}

	q, err := s.getIndexQueue(targetVector)
	if err != nil {
		s.index.logger.WithError(err).Warn("preload queue: queue not found")
		// queue was never initialized, possibly because of a failed shard
		// initialization. No op.
		return nil
	}

	ctx := context.Background()

	maxDocID := s.Counter().Get()

	var batch []common.VectorRecord
	err = s.iterateOnLSMVectors(ctx, from, targetVector, func(id uint64, vector []float32) error {
		if vectorIndex.ContainsNode(id) {
			return nil
		}
		if len(vector) == 0 {
			return nil
		}

		rec := common.VectorRecord{
			ID:     id,
			Vector: vector,
		}
		counter++

		batch = append(batch, rec)

		if len(batch) < 1000 {
			return nil
		}

		err = q.Insert(ctx, batch...)
		if err != nil {
			return err
		}

		batch = batch[:0]
		return nil
	})
	if err != nil {
		return errors.Wrap(err, "iterate on LSM")
	}

	if len(batch) > 0 {
		err = q.Insert(ctx, batch...)
		if err != nil {
			return errors.Wrap(err, "insert batch")
		}
	}

	s.index.logger.
		WithField("last_stored_id", maxDocID).
		WithField("count", counter).
		WithField("took", time.Since(start)).
		WithField("shard_id", s.ID()).
		WithField("target_vector", targetVector).
		Info("enqueued vectors from LSM store")

	return nil
}

func (s *Shard) iterateOnLSMVectors(ctx context.Context, fromID uint64, targetVector string, fn func(id uint64, vector []float32) error) error {
	maxDocID := s.Counter().Get()
	bucket := s.Store().Bucket(helpers.ObjectsBucketLSM)

	buf := make([]byte, 8)
	for i := fromID; i < maxDocID; i++ {
		if ctx.Err() != nil {
			return ctx.Err()
		}

		binary.LittleEndian.PutUint64(buf, i)

		v, err := bucket.GetBySecondary(0, buf)
		if err != nil {
			return errors.Wrap(err, "get last indexed object")
		}
		if v == nil {
			continue
		}
		obj, err := storobj.FromBinary(v)
		if err != nil {
			return errors.Wrap(err, "unmarshal last indexed object")
		}
		id := obj.DocID

		var vector []float32
		if targetVector == "" {
			vector = obj.Vector
		} else {
			if len(obj.Vectors) > 0 {
				vector = obj.Vectors[targetVector]
			}
		}

		err = fn(id, vector)
		if err != nil {
			return err
		}
	}

	return nil
}

// RepairIndex ensures the vector index is consistent with the LSM store.
// It goes through the LSM store and enqueues any unindexed vector, and
// it also removes any indexed vector that is not in the LSM store.
// It it safe to call or interrupt this method at any time.
// If ASYNC_INDEXING is disabled, it's a no-op.
func (s *Shard) RepairIndex(ctx context.Context, targetVector string) error {
	if !asyncEnabled() {
		return nil
	}

	start := time.Now()

	vectorIndex, err := s.getVectorIndex(targetVector)
	if err != nil {
		s.index.logger.Warn("repair index: vector index not found")
		// shard was never initialized, possibly because of a failed shard
		// initialization. No op.
		return nil
	}

	// if it's HNSW, trigger a tombstone cleanup
	if hnsw.IsHNSWIndex(vectorIndex) {
		err := hnsw.AsHNSWIndex(vectorIndex).CleanUpTombstonedNodes(func() bool {
			return ctx.Err() != nil
		})
		if err != nil {
			return errors.Wrap(err, "clean up tombstoned nodes")
		}
	}

	q, err := s.getIndexQueue(targetVector)
	if err != nil {
		s.index.logger.WithError(err).Warn("repair index: queue not found")
		// queue was never initialized, possibly because of a failed shard
		// initialization. No op.
		return nil
	}

	maxDocID := s.Counter().Get()

	visited := visited.NewList(int(maxDocID))

	var added, deleted int

	var batch []common.VectorRecord

	// add non-indexed vectors to the queue
	err = s.iterateOnLSMVectors(ctx, 0, targetVector, func(id uint64, vector []float32) error {
		visited.Visit(id)

		if vectorIndex.ContainsNode(id) {
			return nil
		}
		if len(vector) == 0 {
			return nil
		}

		rec := common.VectorRecord{
			ID:     id,
			Vector: vector,
		}
		added++

		batch = append(batch, rec)

		if len(batch) < 1000 {
			return nil
		}

		err = q.Insert(ctx, batch...)
		if err != nil {
			return err
		}

		batch = batch[:0]
		return nil
	})
	if err != nil {
		return errors.Wrap(err, "iterate on LSM")
	}

	if len(batch) > 0 {
		err = q.Insert(ctx, batch...)
		if err != nil {
			return errors.Wrap(err, "insert batch")
		}
	}

	// if no nodes were visited, it either means the LSM store is empty or
	// there was an uncaught error during the iteration.
	// in any case, we should not touch the index.
	if visited.Len() == 0 {
		s.index.logger.Warn("repair index: empty LSM store")
		return nil
	}

	// remove any indexed vector that is not in the LSM store
	vectorIndex.Iterate(func(id uint64) bool {
		if visited.Visited(id) {
			return true
		}

		deleted++
		err := vectorIndex.Delete(id)
		if err != nil {
			s.index.logger.WithError(err).WithField("id", id).Warn("delete vector from queue")
		}

		return true
	})

	s.index.logger.
		WithField("added", added).
		WithField("deleted", deleted).
		WithField("shard_id", s.ID()).
		WithField("took", time.Since(start)).
		WithField("target_vector", targetVector).
		Info("repaired vector index")

	return nil
}
