//  Copyright (c) 2017 Couchbase, Inc.
//
// Licensed 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 scorch

import (
	"context"
	"errors"
	"fmt"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/RoaringBitmap/roaring/v2"
	"github.com/blevesearch/bleve/v2/index/scorch/mergeplan"
	"github.com/blevesearch/bleve/v2/util"
	segment "github.com/blevesearch/scorch_segment_api/v2"
)

const merger = "merger"

// used in the context of mergerCtrl to provide a way to verify
// the completion of a merge operation
const mergeDoneKey = "mergeDone"

type mergeDoneChan chan error

// used in the context of mergerCtrl to provide a way to use
// a custom merge plan instead of the one generated by the
// default merge planner
const mergePlanFuncKey = "mergePlanFunc"

type mergePlanFunc func(*IndexSnapshot) (*mergeplan.MergePlan, error)

func (s *Scorch) mergerLoop() {
	defer func() {
		if r := recover(); r != nil {
			s.fireAsyncError(NewScorchError(
				merger,
				fmt.Sprintf("panic: %v, path: %s", r, s.path),
				ErrAsyncPanic,
			))
		}

		s.asyncTasks.Done()
	}()

	var lastEpochMergePlanned uint64
	var ctrlMsg *mergerCtrl
	ctrlMsgDflt := &mergerCtrl{ctx: context.Background(),
		options: s.mergePlannerOptions,
		doneCh:  nil}

OUTER:
	for {
		atomic.AddUint64(&s.stats.TotFileMergeLoopBeg, 1)

		select {
		case <-s.closeCh:
			break OUTER

		default:
			// check to see if there is a new snapshot to persist
			s.rootLock.Lock()
			ourSnapshot := s.root
			ourSnapshot.AddRef()
			atomic.StoreUint64(&s.iStats.mergeSnapshotSize, uint64(ourSnapshot.Size()))
			atomic.StoreUint64(&s.iStats.mergeEpoch, ourSnapshot.epoch)
			s.rootLock.Unlock()

			if ctrlMsg == nil && ourSnapshot.epoch != lastEpochMergePlanned {
				ctrlMsg = ctrlMsgDflt
			}
			if ctrlMsg != nil {
				continueMerge := s.fireEvent(EventKindPreMergeCheck, 0)
				// The default, if there's no handler, is to continue the merge.
				if !continueMerge {
					// If it's decided that this merge can't take place now,
					// begin the merge process all over again.
					// Retry instead of blocking/waiting here since a long wait
					// can result in more segments introduced i.e. s.root will
					// be updated.

					// decrement the ref count since its no longer needed in this
					// iteration
					_ = ourSnapshot.DecRef()
					continue OUTER
				}

				// lets get started
				startTime := time.Now()
				err := s.planMergeAtSnapshot(ctrlMsg, ourSnapshot)
				if err != nil {
					atomic.StoreUint64(&s.iStats.mergeEpoch, 0)
					if err == segment.ErrClosed {
						// index has been closed
						_ = ourSnapshot.DecRef()

						// continue the workloop on a user triggered cancel
						if ctrlMsg.doneCh != nil {
							close(ctrlMsg.doneCh)
							ctrlMsg = nil
							continue OUTER
						}

						// exit the workloop on index closure
						ctrlMsg = nil
						break OUTER
					}

					s.fireAsyncError(NewScorchError(
						merger,
						fmt.Sprintf("merging err: %v", err),
						ErrPersist,
					))
					_ = ourSnapshot.DecRef()
					atomic.AddUint64(&s.stats.TotFileMergeLoopErr, 1)
					continue OUTER
				}

				if ctrlMsg.doneCh != nil {
					close(ctrlMsg.doneCh)
				}
				ctrlMsg = nil

				lastEpochMergePlanned = ourSnapshot.epoch

				atomic.StoreUint64(&s.stats.LastMergedEpoch, ourSnapshot.epoch)

				s.fireEvent(EventKindMergerProgress, time.Since(startTime))
			}
			_ = ourSnapshot.DecRef()

			// tell the persister we're waiting for changes
			// first make a epochWatcher chan
			ew := &epochWatcher{
				epoch:    lastEpochMergePlanned,
				notifyCh: make(notificationChan, 1),
			}

			// give it to the persister
			select {
			case <-s.closeCh:
				break OUTER
			case s.persisterNotifier <- ew:
			case ctrlMsg = <-s.forceMergeRequestCh:
				continue OUTER
			}

			// now wait for persister (but also detect close)
			select {
			case <-s.closeCh:
				break OUTER
			case <-ew.notifyCh:
			case ctrlMsg = <-s.forceMergeRequestCh:
			}
		}

		atomic.AddUint64(&s.stats.TotFileMergeLoopEnd, 1)
	}
}

type mergerCtrl struct {
	ctx     context.Context
	options *mergeplan.MergePlanOptions
	doneCh  chan struct{}
}

// ForceMerge helps users trigger a merge operation on
// an online scorch index.
func (s *Scorch) ForceMerge(ctx context.Context,
	mo *mergeplan.MergePlanOptions) error {
	// check whether force merge is already under processing
	s.rootLock.Lock()
	if atomic.LoadUint64(&s.stats.TotFileMergeForceOpsStarted) >
		atomic.LoadUint64(&s.stats.TotFileMergeForceOpsCompleted) {
		s.rootLock.Unlock()
		return fmt.Errorf("force merge already in progress")
	}

	atomic.AddUint64(&s.stats.TotFileMergeForceOpsStarted, 1)
	s.rootLock.Unlock()

	if mo != nil {
		err := mergeplan.ValidateMergePlannerOptions(mo)
		if err != nil {
			return err
		}
	} else {
		// assume the default single segment merge policy
		mo = &mergeplan.SingleSegmentMergePlanOptions
	}
	msg := &mergerCtrl{options: mo,
		doneCh: make(chan struct{}),
		ctx:    ctx,
	}

	// request the merger perform a force merge
	select {
	case s.forceMergeRequestCh <- msg:
	case <-s.closeCh:
		return nil
	}

	// wait for the force merge operation completion
	select {
	case <-msg.doneCh:
		atomic.AddUint64(&s.stats.TotFileMergeForceOpsCompleted, 1)
	case <-s.closeCh:
	}

	return nil
}

func (s *Scorch) parseMergePlannerOptions(po *persisterOptions) (*mergeplan.MergePlanOptions,
	error) {
	mergePlannerOptions := mergeplan.DefaultMergePlanOptions
	// by default use the MaxSizeInMemoryMergePerWorker from the persister option
	// as the FloorSegmentFileSize for the merge planner which would be the
	// first tier size in the planning. If the value is 0, then we don't use the
	// file size in the planning.
	mergePlannerOptions.FloorSegmentFileSize = int64(po.MaxSizeInMemoryMergePerWorker)

	if v, ok := s.config["scorchMergePlanOptions"]; ok {
		b, err := util.MarshalJSON(v)
		if err != nil {
			return &mergePlannerOptions, err
		}

		err = util.UnmarshalJSON(b, &mergePlannerOptions)
		if err != nil {
			return &mergePlannerOptions, err
		}

		err = mergeplan.ValidateMergePlannerOptions(&mergePlannerOptions)
		if err != nil {
			return nil, err
		}
	}
	return &mergePlannerOptions, nil
}

type closeChWrapper struct {
	ch1      chan struct{}
	ctx      context.Context
	closeCh  chan struct{}
	cancelCh chan struct{}
}

func newCloseChWrapper(ch1 chan struct{},
	ctx context.Context) *closeChWrapper {
	return &closeChWrapper{
		ch1:      ch1,
		ctx:      ctx,
		closeCh:  make(chan struct{}),
		cancelCh: make(chan struct{}),
	}
}

func (w *closeChWrapper) close() {
	close(w.closeCh)
}

func (w *closeChWrapper) listen() {
	select {
	case <-w.ch1:
		close(w.cancelCh)
	case <-w.ctx.Done():
		close(w.cancelCh)
	case <-w.closeCh:
	}
}

// mergeBatch represents a batch of segments to be merged together,
// along with the new segment that will be created as a result of the merge.
type mergeBatch struct {
	snapshots []*SegmentSnapshot
	filenames []string
	segments  []segment.Segment
	drops     []*roaring.Bitmap

	new         segment.Segment
	newID       uint64
	newDocNums  [][]uint64
	newFilename string
	newTime     uint64
}

// planMergeAtSnapshot plans and executes the merge operations for a given snapshot
// if there is a custom merge plan function provided, it uses that to get the merge plan
// otherwise, it builds the merge plan using the default planner and executes the merge tasks in the plan.
func (s *Scorch) planMergeAtSnapshot(ctrlMsg *mergerCtrl, ourSnapshot *IndexSnapshot) error {
	var mergePlan *mergeplan.MergePlan
	// if a merge plan function is provided in the context, use it to get the merge plan
	if mergePlanFunc, ok := ctrlMsg.ctx.Value(mergePlanFuncKey).(mergePlanFunc); ok {
		var err error
		mergePlan, err = mergePlanFunc(ourSnapshot)
		if err != nil {
			atomic.AddUint64(&s.stats.TotFileMergePlanErr, 1)
			return fmt.Errorf("merge planning err: %v", err)
		}
	}

	// default to making a merge plan if a custom one is not provided
	if mergePlan == nil || len(mergePlan.Tasks) == 0 {
		// build list of persisted segments in this snapshot
		var onlyPersistedSnapshots []mergeplan.Segment
		for _, segmentSnapshot := range ourSnapshot.segment {
			if _, ok := segmentSnapshot.segment.(segment.PersistedSegment); ok {
				onlyPersistedSnapshots = append(onlyPersistedSnapshots, segmentSnapshot)
			}
		}

		atomic.AddUint64(&s.stats.TotFileMergePlan, 1)

		// give this list to the planner
		var err error
		mergePlan, err = mergeplan.Plan(onlyPersistedSnapshots, ctrlMsg.options)
		if err != nil {
			atomic.AddUint64(&s.stats.TotFileMergePlanErr, 1)
			return fmt.Errorf("merge planning err: %v", err)
		}

		if mergePlan == nil || len(mergePlan.Tasks) == 0 {
			// nothing to do
			atomic.AddUint64(&s.stats.TotFileMergePlanNone, 1)
			return nil
		}
	}

	atomic.AddUint64(&s.stats.TotFileMergePlanOk, 1)
	atomic.AddUint64(&s.stats.TotFileMergePlanTasks, uint64(len(mergePlan.Tasks)))

	cw := newCloseChWrapper(s.closeCh, ctrlMsg.ctx)
	defer cw.close()
	go cw.listen()

	var err error
	defer func() {
		// send error to done channel if present
		if done, ok := cw.ctx.Value(mergeDoneKey).(chan error); ok {
			done <- err
		}
	}()

	numBatches := len(mergePlan.Tasks)
	mergeBatches := make([]*mergeBatch, numBatches)
	defer func() {
		for batchID := 0; batchID < numBatches; batchID++ {
			batch := mergeBatches[batchID]
			if batch == nil {
				continue
			}
			if err == nil {
				s.rootLock.Lock()
				for _, filename := range batch.filenames {
					delete(s.ineligibleForRemoval, filename)
				}
				s.rootLock.Unlock()
			} else {
				if batch.new != nil {
					_ = batch.new.Close()
				}
				s.unmarkIneligibleForRemoval(batch.newFilename)
			}
		}
	}()

	numMergedSegments := 0
	for batchID := 0; batchID < numBatches; batchID++ {
		task := mergePlan.Tasks[batchID]
		if len(task.Segments) == 0 {
			atomic.AddUint64(&s.stats.TotFileMergePlanTasksSegmentsEmpty, 1)
		}
		atomic.AddUint64(&s.stats.TotFileMergePlanTasksSegments, uint64(len(task.Segments)))

		newSegmentID := atomic.AddUint64(&s.nextSegmentID, 1)
		batch := &mergeBatch{
			snapshots: make([]*SegmentSnapshot, 0, len(task.Segments)),
			filenames: make([]string, 0, len(task.Segments)),
			segments:  make([]segment.Segment, 0, len(task.Segments)),
			drops:     make([]*roaring.Bitmap, 0, len(task.Segments)),

			new:         nil,
			newID:       newSegmentID,
			newDocNums:  nil,
			newFilename: zapFileName(newSegmentID),
			newTime:     0,
		}
		mergeBatches[batchID] = batch

		for _, planSegment := range task.Segments {
			if segSnapshot, ok := planSegment.(*SegmentSnapshot); ok {
				if persistedSeg, ok := segSnapshot.segment.(segment.PersistedSegment); ok {
					if segSnapshot.LiveSize() == 0 {
						atomic.AddUint64(&s.stats.TotFileMergeSegmentsEmpty, 1)
					} else {
						batch.snapshots = append(batch.snapshots, segSnapshot)
						batch.segments = append(batch.segments, segSnapshot.segment)
						batch.drops = append(batch.drops, segSnapshot.deleted)
					}
					// track the files getting merged for unsetting the
					// removal ineligibility. This helps to unflip files
					// even with fast merger, slow persister work flows.
					path := persistedSeg.Path()
					batch.filenames = append(batch.filenames,
						strings.TrimPrefix(path, s.path+string(os.PathSeparator)))
				}
			}
		}

		if len(batch.segments) == 0 {
			continue
		}

		numMergedSegments += len(batch.snapshots)
		s.markIneligibleForRemoval(batch.newFilename)
		path := s.path + string(os.PathSeparator) + batch.newFilename

		prevBytesReadTotal := cumulateBytesRead(batch.segments)
		fileMergeZapStartTime := time.Now()
		atomic.AddUint64(&s.stats.TotFileMergeZapBeg, 1)
		batch.newDocNums, _, err = s.segPlugin.MergeUsing(batch.segments, batch.drops, path,
			cw.cancelCh, s, s.segmentConfig)
		atomic.AddUint64(&s.stats.TotFileMergeZapEnd, 1)
		batch.newTime = uint64(time.Since(fileMergeZapStartTime))
		atomic.AddUint64(&s.stats.TotFileMergeZapTime, batch.newTime)
		if atomic.LoadUint64(&s.stats.MaxFileMergeZapTime) < batch.newTime {
			atomic.StoreUint64(&s.stats.MaxFileMergeZapTime, batch.newTime)
		}
		if err != nil {
			atomic.AddUint64(&s.stats.TotFileMergePlanTasksErr, 1)
			return err
		}

		batch.new, err = s.segPlugin.OpenUsing(path, s.segmentConfig)
		if err != nil {
			atomic.AddUint64(&s.stats.TotFileMergePlanTasksErr, 1)
			return err
		}

		totalBytesRead := batch.new.BytesRead() + prevBytesReadTotal
		batch.new.ResetBytesRead(totalBytesRead)
		atomic.AddUint64(&s.stats.TotFileMergeSegments, uint64(len(batch.segments)))
	}

	newSegmentIDs := make([]uint64, numBatches)
	newSegments := make([]segment.Segment, numBatches)
	mergedSegHistory := make(map[uint64]*mergedSegmentHistory, numMergedSegments)
	for batchID := 0; batchID < numBatches; batchID++ {
		batch := mergeBatches[batchID]
		newSegmentIDs[batchID] = batch.newID
		newSegments[batchID] = batch.new
		for j, ss := range batch.snapshots {
			mergedSegHistory[ss.id] = &mergedSegmentHistory{
				batchID:      batchID,
				oldNewDocIDs: batch.newDocNums[j],
				oldSegment:   ss,
			}
		}
	}

	sm := &segmentMerge{
		newSegmentIDs:    newSegmentIDs,
		newSegments:      newSegments,
		mergedSegHistory: mergedSegHistory,
		notifyCh:         make(chan *mergeTaskIntroStatus),
		mmaped:           1,
		fileMerge:        true,
	}

	s.fireEvent(EventKindMergeTaskIntroductionStart, 0)
	// give it to the introducer
	select {
	case <-s.closeCh:
		err = segment.ErrClosed
		return err
	case s.merges <- sm:
		atomic.AddUint64(&s.stats.TotFileMergeIntroductions, uint64(numBatches))
	}

	introStartTime := time.Now()
	// it is safe to blockingly wait for the merge introduction
	// here as the introducer is bound to handle the notify channel.
	introStatus := <-sm.notifyCh
	introTime := uint64(time.Since(introStartTime))
	atomic.AddUint64(&s.stats.TotFileMergeZapIntroductionTime, introTime)
	if atomic.LoadUint64(&s.stats.MaxFileMergeZapIntroductionTime) < introTime {
		atomic.StoreUint64(&s.stats.MaxFileMergeZapIntroductionTime, introTime)
	}
	atomic.AddUint64(&s.stats.TotFileMergeIntroductionsDone, uint64(numBatches))

	_ = introStatus.indexSnapshot.DecRef()
	for batchID, skipped := range introStatus.skipped {
		if skipped {
			batch := mergeBatches[batchID]
			if batch.new != nil {
				_ = batch.new.Close()
			}
			s.unmarkIneligibleForRemoval(batch.newFilename)
		}
	}

	atomic.AddUint64(&s.stats.TotFileMergePlanTasksDone, uint64(numBatches))
	s.fireEvent(EventKindMergeTaskIntroduction, 0)
	return nil
}

type mergeTaskIntroStatus struct {
	indexSnapshot *IndexSnapshot
	skipped       []bool
}

// this is important when it comes to introducing multiple merged segments in a
// single introducer channel push. That way there is a check to ensure that the
// file count doesn't explode during the index's lifetime.
type mergedSegmentHistory struct {
	batchID      int
	oldNewDocIDs []uint64
	oldSegment   *SegmentSnapshot
}

type segmentMerge struct {
	newSegmentIDs    []uint64
	newSegments      []segment.Segment
	mergedSegHistory map[uint64]*mergedSegmentHistory
	notifyCh         chan *mergeTaskIntroStatus
	mmaped           uint32
	fileMerge        bool
}

func cumulateBytesRead(sbs []segment.Segment) uint64 {
	var rv uint64
	for _, seg := range sbs {
		rv += seg.BytesRead()
	}
	return rv
}

func (s *Scorch) mergeAndPersistInMemorySegments(flushes []*flushable, po *persisterOptions) (*IndexSnapshot, map[uint64]struct{}, error) {
	atomic.AddUint64(&s.stats.TotMemMergeBeg, 1)

	var err error
	numBatches := len(flushes)
	mergeBatches := make([]*mergeBatch, numBatches)
	defer func() {
		if err != nil {
			for batchID := 0; batchID < numBatches; batchID++ {
				batch := mergeBatches[batchID]
				if batch.new != nil {
					_ = batch.new.Close()
				}
				s.unmarkIneligibleForRemoval(batch.newFilename)
			}
		}
	}()

	var numMergedSegments uint64
	sem := make(chan struct{}, po.NumPersisterWorkers)
	var wg sync.WaitGroup
	var errM sync.Mutex
	var errs []error
	for batchID := 0; batchID < numBatches; batchID++ {
		wg.Add(1)
		sem <- struct{}{}
		go func(batchID int) {
			var err error
			defer func() {
				if err != nil {
					atomic.AddUint64(&s.stats.TotMemMergeErr, 1)
					errM.Lock()
					errs = append(errs, err)
					errM.Unlock()
				}
				<-sem
				wg.Done()
			}()

			flush := flushes[batchID]
			newSegmentID := atomic.AddUint64(&s.nextSegmentID, 1)
			batch := &mergeBatch{
				snapshots: flush.sbsBatchSnapshots,
				filenames: nil,
				segments:  flush.sbsBatch,
				drops:     flush.sbsBatchDrops,

				new:         nil,
				newID:       newSegmentID,
				newDocNums:  nil,
				newFilename: zapFileName(newSegmentID),
				newTime:     0,
			}
			mergeBatches[batchID] = batch

			s.markIneligibleForRemoval(batch.newFilename)
			path := s.path + string(os.PathSeparator) + batch.newFilename

			memMergeZapStartTime := time.Now()
			atomic.AddUint64(&s.stats.TotMemMergeZapBeg, 1)
			batch.newDocNums, _, err = s.segPlugin.MergeUsing(batch.segments, batch.drops, path,
				s.closeCh, s, s.segmentConfig)
			atomic.AddUint64(&s.stats.TotMemMergeZapEnd, 1)
			batch.newTime = uint64(time.Since(memMergeZapStartTime))
			atomic.AddUint64(&s.stats.TotMemMergeZapTime, batch.newTime)
			if err != nil {
				return
			}
			batch.new, err = s.segPlugin.OpenUsing(path, s.segmentConfig)
			if err != nil {
				return
			}
			atomic.AddUint64(&numMergedSegments, uint64(len(batch.segments)))
		}(batchID)
	}
	wg.Wait()
	close(sem)

	var maxZapTime uint64
	for _, batch := range mergeBatches {
		if batch.newTime > maxZapTime {
			maxZapTime = batch.newTime
		}
	}
	if maxZapTime > atomic.LoadUint64(&s.stats.MaxMemMergeZapTime) {
		atomic.StoreUint64(&s.stats.MaxMemMergeZapTime, maxZapTime)
	}
	atomic.AddUint64(&s.stats.TotMemMergeSegments, numMergedSegments)

	if len(errs) > 0 {
		var errf error
		numClosed := 0
		for _, e := range errs {
			if e == segment.ErrClosed {
				numClosed++
			}
			errf = errors.Join(errf, e)
		}
		if numClosed == len(errs) {
			// the index snapshot was closed which will be handled gracefully
			// by retrying the whole merge+flush operation in a later iteration
			// so its safe to early exit the same error.
			err = segment.ErrClosed
		} else {
			err = errf
		}
		return nil, nil, err
	}

	newSegmentIDs := make([]uint64, numBatches)
	newSegments := make([]segment.Segment, numBatches)
	mergedSegHistory := make(map[uint64]*mergedSegmentHistory, numMergedSegments)
	for batchID := 0; batchID < numBatches; batchID++ {
		batch := mergeBatches[batchID]
		newSegmentIDs[batchID] = batch.newID
		newSegments[batchID] = batch.new
		for j, ss := range batch.snapshots {
			mergedSegHistory[ss.id] = &mergedSegmentHistory{
				batchID:      batchID,
				oldNewDocIDs: batch.newDocNums[j],
				oldSegment:   ss,
			}
		}
	}

	sm := &segmentMerge{
		newSegmentIDs:    newSegmentIDs,
		newSegments:      newSegments,
		mergedSegHistory: mergedSegHistory,
		notifyCh:         make(chan *mergeTaskIntroStatus),
		mmaped:           1,
		fileMerge:        false,
	}

	select {
	case <-s.closeCh:
		err = segment.ErrClosed
		return nil, nil, err
	case s.merges <- sm:
		atomic.AddUint64(&s.stats.TotMemMergeIntroductions, uint64(numBatches))
	}

	introStartTime := time.Now()
	// blockingly wait for the introduction to complete
	introStatus := <-sm.notifyCh
	introTime := uint64(time.Since(introStartTime))
	atomic.AddUint64(&s.stats.TotMemMergeZapIntroductionTime, introTime)
	if atomic.LoadUint64(&s.stats.MaxMemMergeZapIntroductionTime) < introTime {
		atomic.StoreUint64(&s.stats.MaxMemMergeZapIntroductionTime, introTime)
	}
	atomic.AddUint64(&s.stats.TotMemMergeIntroductionsDone, uint64(numBatches))
	introducedSnapshot := introStatus.indexSnapshot
	introducedSegmentIDs := make(map[uint64]struct{}, numBatches)
	for batchID, skipped := range introStatus.skipped {
		batch := mergeBatches[batchID]
		if skipped {
			_ = batch.new.Close()
			s.unmarkIneligibleForRemoval(batch.newFilename)
		} else {
			introducedSegmentIDs[batch.newID] = struct{}{}
		}
	}

	atomic.AddUint64(&s.stats.TotMemMergeDone, 1)

	// if we could not introduce all of the newly merged segments,
	// then we cannot persist a snapshot with any merged segments at all
	if len(introducedSegmentIDs) != numBatches {
		_ = introducedSnapshot.DecRef()
		return nil, nil, nil
	}

	return introducedSnapshot, introducedSegmentIDs, nil
}

func (s *Scorch) ReportBytesWritten(bytesWritten uint64) {
	atomic.AddUint64(&s.stats.TotFileMergeWrittenBytes, bytesWritten)
}