package stats

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/cockroachdb/errors"

	"github.com/milvus-io/milvus/internal/flushcommon/metacache"
	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors/shard/policy"
	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors/shard/utils"
	"github.com/milvus-io/milvus/pkg/v3/mlog"
	"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)

type (
	SegmentStats         = utils.SegmentStats
	ModifiedMetrics      = utils.ModifiedMetrics
	SegmentBelongs       = utils.SegmentBelongs
	SyncOperationMetrics = utils.SyncOperationMetrics
	SealOperator         = utils.SealOperator
)

var (
	ErrNotEnoughSpace = errors.New("not enough space")
	ErrTooLargeInsert = errors.New("insert too large")
)

// StatsManager is the manager of stats.
// It manages the insert stats of all segments, used to check if a segment has enough space to insert or should be sealed.
// If there will be a lock contention, we can optimize it by apply lock per segment.
type StatsManager struct {
	mlog.Binder
	worker                     *sealWorker
	mu                         sync.Mutex
	cfg                        statsConfig
	totalStats                 *aggregatedMetrics
	totalFlushSize             uint64
	pchannelStats              map[string]*aggregatedMetrics
	vchannelStats              map[string]*aggregatedMetrics
	segmentStats               map[int64]*SegmentStats  // map[SegmentID]SegmentStats
	segmentIndex               map[int64]SegmentBelongs // map[SegmentID]channels
	segmentFlushSourceModes    map[int64]metacache.FlushSourceMode
	pchannelIndex              map[string]map[int64]struct{} // map[PChannel]SegmentID
	growingL1SegmentsByChannel map[channelKey]map[int64]struct{}
	segmentDeletePressures     map[int64]deletePressure // map[SegmentID]aggregated delete pressure
	sealOperators              map[string]SealOperator
	metricHelper               *metricsHelper
}

type channelKey struct {
	pchannel string
	vchannel string
}

type deletePressure struct {
	rows  uint64
	bytes uint64
}

// sealSegmentIDWithPolicy is the struct that contains the segment ID and the seal policy.
type sealSegmentIDWithPolicy struct {
	segmentID  int64
	sealPolicy policy.SealPolicy
}

// NewStatsManager creates a new stats manager.
func NewStatsManager() *StatsManager {
	cfg := newStatsConfig()
	if err := cfg.Validate(); err != nil {
		panic(err)
	}
	m := &StatsManager{
		mu:                         sync.Mutex{},
		cfg:                        cfg,
		totalStats:                 newAggregatedMetrics(),
		pchannelStats:              make(map[string]*aggregatedMetrics),
		vchannelStats:              make(map[string]*aggregatedMetrics),
		segmentStats:               make(map[int64]*SegmentStats),
		segmentIndex:               make(map[int64]SegmentBelongs),
		segmentFlushSourceModes:    make(map[int64]metacache.FlushSourceMode),
		pchannelIndex:              make(map[string]map[int64]struct{}),
		growingL1SegmentsByChannel: make(map[channelKey]map[int64]struct{}),
		segmentDeletePressures:     make(map[int64]deletePressure),
		sealOperators:              make(map[string]SealOperator),
		metricHelper:               newMetricsHelper(),
	}
	m.worker = newSealWorker(m)
	m.metricHelper.ObserveFlushPressureBytesUpdate(m.totalFlushSize)
	go m.worker.loop()
	return m
}

// RegisterSealOperator registers a seal operator and current growing segments related to the seal operator.
// It will perform an atomic operation to register the seal operator and segments into the manager.
func (m *StatsManager) RegisterSealOperator(sealOperator SealOperator, belongs []SegmentBelongs, stats []*SegmentStats) {
	m.registerSealOperator(sealOperator, belongs, stats)
	m.notifyIfTotalGrowingBytesOverHWM()
}

func (m *StatsManager) registerSealOperator(sealOperator SealOperator, belongs []SegmentBelongs, stats []*SegmentStats) {
	if len(belongs) != len(stats) {
		panic("register a seal operator with different length of belongs and stats, critical bug")
	}
	m.mu.Lock()
	defer m.mu.Unlock()

	if _, ok := m.sealOperators[sealOperator.Channel().Name]; ok {
		panic(fmt.Sprintf("register a seal operator %s that already exist, critical bug", sealOperator.Channel().Name))
	}
	m.sealOperators[sealOperator.Channel().Name] = sealOperator
	for i := range belongs {
		m.registerNewGrowingSegment(belongs[i], stats[i])
	}
}

// UnregisterSealOperator unregisters the seal operator and all segments related to it.
func (m *StatsManager) UnregisterSealOperator(sealOperator SealOperator) {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Remove the seal operator from the map
	if _, ok := m.sealOperators[sealOperator.Channel().Name]; !ok {
		panic(fmt.Sprintf("unregister a seal operator %s that not exist, critical bug", sealOperator.Channel().Name))
	}
	delete(m.sealOperators, sealOperator.Channel().Name)

	// Unregister all segments on the seal operator
	m.unregisterAllStatsOnPChannel(sealOperator.Channel().Name)
}

// unregisterAllStatsOnPChannel unregisters all stats on pchannel.
func (m *StatsManager) unregisterAllStatsOnPChannel(pchannel string) int {
	segmentIDs, ok := m.pchannelIndex[pchannel]
	if !ok {
		return 0
	}
	for segmentID := range segmentIDs {
		m.unregisterSealedSegment(segmentID)
	}
	return len(segmentIDs)
}

// RegisterNewGrowingSegment registers a new growing segment.
// delegate the stats management to stats manager.
// It must be called after RegisterSealOperator and before UnregisterSealOperator.
func (m *StatsManager) RegisterNewGrowingSegment(belongs SegmentBelongs, stats *SegmentStats) {
	m.registerNewGrowingSegmentWithMutex(belongs, stats)
	m.notifyIfTotalGrowingBytesOverHWM()
}

func (m *StatsManager) registerNewGrowingSegmentWithMutex(belongs SegmentBelongs, stats *SegmentStats) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if _, ok := m.sealOperators[belongs.PChannel]; !ok {
		panic(fmt.Sprintf("register a segment %+v that seal operator is not exist, critical bug", belongs))
	}
	m.registerNewGrowingSegment(belongs, stats)
}

// registerNewGrowingSegment registers a new growing segment.
// delegate the stats management to stats manager.
func (m *StatsManager) registerNewGrowingSegment(belongs SegmentBelongs, stats *SegmentStats) {
	segmentID := belongs.SegmentID
	if _, ok := m.segmentStats[segmentID]; ok {
		panic(fmt.Sprintf("register a segment %d that already exist, critical bug", segmentID))
	}
	if stats.RuntimeFlushSize == 0 {
		stats.RuntimeFlushSize = stats.Modified.BinarySize
	}

	m.segmentStats[segmentID] = stats
	m.segmentIndex[segmentID] = belongs
	if _, ok := m.pchannelIndex[belongs.PChannel]; !ok {
		m.pchannelIndex[belongs.PChannel] = make(map[int64]struct{})
	}
	m.pchannelIndex[belongs.PChannel][segmentID] = struct{}{}
	if stats.Level == datapb.SegmentLevel_L1 {
		key := channelKey{pchannel: belongs.PChannel, vchannel: belongs.VChannel}
		if _, ok := m.growingL1SegmentsByChannel[key]; !ok {
			m.growingL1SegmentsByChannel[key] = make(map[int64]struct{})
		}
		m.growingL1SegmentsByChannel[key][segmentID] = struct{}{}
	}
	m.totalStats.Collect(stats.Level, stats.Modified)
	m.totalFlushSize = utils.SaturatingAddUint64(m.totalFlushSize, stats.FlushSize())
	m.metricHelper.ObserveFlushPressureBytesUpdate(m.totalFlushSize)
	if _, ok := m.pchannelStats[belongs.PChannel]; !ok {
		m.pchannelStats[belongs.PChannel] = newAggregatedMetrics()
	}
	m.pchannelStats[belongs.PChannel].Collect(stats.Level, stats.Modified)

	if _, ok := m.vchannelStats[belongs.VChannel]; !ok {
		m.vchannelStats[belongs.VChannel] = newAggregatedMetrics()
	}
	m.vchannelStats[belongs.VChannel].Collect(stats.Level, stats.Modified)

	m.metricHelper.ObservePChannelBytesUpdate(belongs.PChannel, m.pchannelStats[belongs.PChannel])
}

// AllocRows alloc number of rows on current segment.
// AllocRows will check if the segment has enough space to insert.
// Must be called after RegisterGrowingSegment and before UnregisterGrowingSegment.
func (m *StatsManager) AllocRows(segmentID int64, insert ModifiedMetrics, runtimeFlushSize ...uint64) error {
	if insert.Rows == 0 || insert.BinarySize == 0 {
		panic(fmt.Sprintf("insert rows or binary size cannot be 0, rows: %d, binary: %d", insert.Rows, insert.BinarySize))
	}

	shouldBeSealed, err := m.allocRows(segmentID, insert, normalizeRuntimeFlushSize(insert, runtimeFlushSize...))
	if shouldBeSealed {
		m.worker.NotifySealSegment(segmentID, policy.PolicyCapacity())
	}
	if err != nil {
		return err
	}
	m.notifyIfTotalGrowingBytesOverHWM()
	return nil
}

func normalizeRuntimeFlushSize(insert ModifiedMetrics, runtimeFlushSize ...uint64) uint64 {
	if len(runtimeFlushSize) == 0 || runtimeFlushSize[0] < insert.BinarySize {
		return insert.BinarySize
	}
	return runtimeFlushSize[0]
}

// allocRows allocates number of rows on current segment.
func (m *StatsManager) allocRows(segmentID int64, insert ModifiedMetrics, runtimeFlushSize uint64) (bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Must be exist, otherwise it's a bug.
	info, ok := m.segmentIndex[segmentID]
	if !ok {
		panic(fmt.Sprintf("alloc rows on a segment %d that not exist", segmentID))
	}
	stat := m.segmentStats[segmentID]
	inserted := stat.AllocRows(insert)

	// update the total stats if inserted.
	if inserted {
		if m.segmentFlushSourceModes[segmentID] == metacache.FlushSourceWriteBuffer {
			runtimeFlushSize = insert.BinarySize
		}
		stat.AllocRuntimeFlushSize(runtimeFlushSize)
		m.totalFlushSize = utils.SaturatingAddUint64(m.totalFlushSize, runtimeFlushSize)
		m.metricHelper.ObserveFlushPressureBytesUpdate(m.totalFlushSize)
		m.totalStats.Collect(stat.Level, insert)
		if _, ok := m.pchannelStats[info.PChannel]; !ok {
			m.pchannelStats[info.PChannel] = newAggregatedMetrics()
		}
		m.pchannelStats[info.PChannel].Collect(stat.Level, insert)
		if _, ok := m.vchannelStats[info.VChannel]; !ok {
			m.vchannelStats[info.VChannel] = newAggregatedMetrics()
		}
		m.vchannelStats[info.VChannel].Collect(stat.Level, insert)

		m.metricHelper.ObservePChannelBytesUpdate(info.PChannel, m.pchannelStats[info.PChannel])
		return stat.ShouldBeSealed(), nil
	}
	if stat.IsEmpty() {
		return false, ErrTooLargeInsert
	}
	return stat.ShouldBeSealed(), ErrNotEnoughSpace
}

// UpdateFlushSourceMode records the segment-level sticky flush source chosen by writebuffer.
func (m *StatsManager) UpdateFlushSourceMode(segmentID int64, mode metacache.FlushSourceMode) {
	if mode != metacache.FlushSourceWriteBuffer && mode != metacache.FlushSourceGrowing {
		return
	}

	m.mu.Lock()
	defer m.mu.Unlock()

	stat, ok := m.segmentStats[segmentID]
	if !ok {
		return
	}
	current := m.segmentFlushSourceModes[segmentID]
	if current != metacache.FlushSourceUnknown && current != mode {
		return
	}
	m.segmentFlushSourceModes[segmentID] = mode
	if mode == metacache.FlushSourceWriteBuffer {
		oldFlushSize := stat.FlushSize()
		stat.RuntimeFlushSize = stat.Modified.BinarySize
		newFlushSize := stat.FlushSize()
		if oldFlushSize > newFlushSize {
			m.totalFlushSize = utils.SaturatingSubUint64(m.totalFlushSize, oldFlushSize-newFlushSize)
		} else if newFlushSize > oldFlushSize {
			m.totalFlushSize = utils.SaturatingAddUint64(m.totalFlushSize, newFlushSize-oldFlushSize)
		}
		m.metricHelper.ObserveFlushPressureBytesUpdate(m.totalFlushSize)
	}
}

// RecordDelete records local delete metrics on matching growing L1 segments.
func (m *StatsManager) RecordDelete(pchannel string, vchannel string, timeTick uint64, rows uint64, bytes uint64) {
	if pchannel == "" || vchannel == "" || timeTick == 0 || (rows == 0 && bytes == 0) {
		return
	}

	m.mu.Lock()
	defer m.mu.Unlock()

	key := channelKey{pchannel: pchannel, vchannel: vchannel}
	for segmentID := range m.growingL1SegmentsByChannel[key] {
		stat := m.segmentStats[segmentID]
		if stat.CreateSegmentTimeTick == 0 || timeTick < stat.CreateSegmentTimeTick {
			continue
		}
		pressure := m.segmentDeletePressures[segmentID]
		pressure.rows += rows
		pressure.bytes += bytes
		m.segmentDeletePressures[segmentID] = pressure
	}
}

// notifyIfTotalGrowingBytesOverHWM notifies if the total bytes is over the high water mark.
func (m *StatsManager) notifyIfTotalGrowingBytesOverHWM() {
	m.mu.Lock()
	size := m.totalFlushSize
	notify := size > uint64(m.cfg.growingBytesHWM)
	m.mu.Unlock()

	if notify {
		m.worker.NotifyGrowingBytes(size)
	}
}

// GetStatsOfSegment gets the stats of segment.
func (m *StatsManager) GetStatsOfSegment(segmentID int64) *SegmentStats {
	m.mu.Lock()
	defer m.mu.Unlock()

	return m.segmentStats[segmentID].Copy()
}

// getSealOperator gets the seal operator of the segment.
func (m *StatsManager) getSealOperator(segmentID int64) (SegmentBelongs, *SegmentStats, SealOperator, bool) {
	m.mu.Lock()
	defer m.mu.Unlock()

	belongs, ok := m.segmentIndex[segmentID]
	if !ok {
		return belongs, nil, nil, false
	}
	stats, ok := m.segmentStats[segmentID]
	if !ok {
		panic(fmt.Sprintf("stats of segment %d that not exist, critical bug", segmentID))
	}
	sealOperator, ok := m.sealOperators[belongs.PChannel]
	if !ok {
		panic(fmt.Sprintf("seal operator of segment %d that not exist, critical bug", segmentID))
	}
	return belongs, stats.Copy(), sealOperator, true
}

// UpdateOnSync updates the stats of segment on sync.
// It's an async update operation, so it's not necessary to do success.
func (m *StatsManager) UpdateOnSync(segmentID int64, syncMetric SyncOperationMetrics) {
	m.mu.Lock()
	if _, ok := m.segmentIndex[segmentID]; !ok {
		// UpdateOnSync is called asynchronously, so we need to check if the segment is still exist.
		m.mu.Unlock()
		return
	}
	m.segmentStats[segmentID].UpdateOnSync(syncMetric)
	limit := uint64(m.cfg.maxBinlogFileNum)
	notify := m.segmentStats[segmentID].BinLogCounter > limit
	m.mu.Unlock()

	// Trigger seal if the binlog file number reach the limit.
	if notify {
		m.worker.NotifySealSegment(segmentID, policy.PolicyBinlogNumber(limit))
	}
}

// UnregisterSealedSegment unregisters the sealed segment.
func (m *StatsManager) UnregisterSealedSegment(segmentID int64) *SegmentStats {
	m.mu.Lock()
	defer m.mu.Unlock()

	return m.unregisterSealedSegment(segmentID)
}

func (m *StatsManager) unregisterSealedSegment(segmentID int64) *SegmentStats {
	// Must be exist, otherwise it's a bug.
	info, ok := m.segmentIndex[segmentID]
	if !ok {
		panic(fmt.Sprintf("unregister a segment %d that not exist, critical bug", segmentID))
	}

	stats := m.segmentStats[segmentID]

	m.totalStats.Subtract(stats.Level, stats.Modified)
	m.totalFlushSize = utils.SaturatingSubUint64(m.totalFlushSize, stats.FlushSize())
	m.metricHelper.ObserveFlushPressureBytesUpdate(m.totalFlushSize)
	delete(m.segmentStats, segmentID)
	delete(m.segmentIndex, segmentID)
	delete(m.segmentFlushSourceModes, segmentID)
	delete(m.segmentDeletePressures, segmentID)
	if stats.Level == datapb.SegmentLevel_L1 {
		key := channelKey{pchannel: info.PChannel, vchannel: info.VChannel}
		delete(m.growingL1SegmentsByChannel[key], segmentID)
		if len(m.growingL1SegmentsByChannel[key]) == 0 {
			delete(m.growingL1SegmentsByChannel, key)
		}
	}
	if _, ok := m.pchannelIndex[info.PChannel]; ok {
		delete(m.pchannelIndex[info.PChannel], segmentID)
		if len(m.pchannelIndex[info.PChannel]) == 0 {
			delete(m.pchannelIndex, info.PChannel)
		}
	}

	if _, ok := m.pchannelStats[info.PChannel]; ok {
		m.pchannelStats[info.PChannel].Subtract(stats.Level, stats.Modified)
		m.metricHelper.ObservePChannelBytesUpdate(info.PChannel, m.pchannelStats[info.PChannel])
		if m.pchannelStats[info.PChannel].IsZero() {
			// If the binary size is 0, it means the segment is empty, we can delete it.
			delete(m.pchannelStats, info.PChannel)
		}
	}
	if _, ok := m.vchannelStats[info.VChannel]; ok {
		m.vchannelStats[info.VChannel].Subtract(stats.Level, stats.Modified)
		if m.vchannelStats[info.VChannel].IsZero() {
			delete(m.vchannelStats, info.VChannel)
		}
	}
	return stats
}

// selectSegmensWithTimePolicy selects segments with time policy.
func (m *StatsManager) selectSegmentsWithTimePolicy() map[int64]policy.SealPolicy {
	m.mu.Lock()
	defer m.mu.Unlock()

	now := time.Now()
	sealSegmentIDs := make(map[int64]policy.SealPolicy, 0)
	for segmentID, stat := range m.segmentStats {
		switch stat.Level {
		case datapb.SegmentLevel_L1:
			if now.Sub(stat.CreateTime) > m.cfg.l1MaxLifetime {
				sealSegmentIDs[segmentID] = policy.PolicyLifetime(m.cfg.l1MaxLifetime)
				continue
			}
			if stat.Modified.BinarySize > uint64(m.cfg.l1MinSizeFromIdleTime) && now.Sub(stat.LastModifiedTime) > m.cfg.l1MaxIdleTime {
				sealSegmentIDs[segmentID] = policy.PolicyIdle(m.cfg.l1MaxIdleTime, uint64(m.cfg.l1MinSizeFromIdleTime))
				continue
			}
		case datapb.SegmentLevel_L0:
			if now.Sub(stat.CreateTime) > m.cfg.l0MaxLifetime {
				sealSegmentIDs[segmentID] = policy.PolicyLifetime(m.cfg.l0MaxLifetime)
				continue
			}
		}
	}
	return sealSegmentIDs
}

// selectSegmentsWithBlockingL0Policy selects the earliest L1 growing segment per vchannel
// when local deletes blocked by that segment reach the configured threshold.
func (m *StatsManager) selectSegmentsWithBlockingL0Policy() map[int64]policy.SealPolicy {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.cfg.blockingL0EntryNum < 0 && m.cfg.blockingL0SizeBytes < 0 {
		return nil
	}

	type pressureScope struct {
		pchannel string
		vchannel string
	}
	type earliestSegment struct {
		segmentID int64
		timeTick  uint64
	}

	earliestByScope := make(map[pressureScope]earliestSegment)
	for segmentID, stat := range m.segmentStats {
		if stat.Level != datapb.SegmentLevel_L1 || stat.CreateSegmentTimeTick == 0 {
			continue
		}
		belongs := m.segmentIndex[segmentID]
		scope := pressureScope{
			pchannel: belongs.PChannel,
			vchannel: belongs.VChannel,
		}
		earliest, ok := earliestByScope[scope]
		if !ok || stat.CreateSegmentTimeTick < earliest.timeTick {
			earliestByScope[scope] = earliestSegment{
				segmentID: segmentID,
				timeTick:  stat.CreateSegmentTimeTick,
			}
		}
	}

	sealSegmentIDs := make(map[int64]policy.SealPolicy)
	for _, earliest := range earliestByScope {
		pressure := m.segmentDeletePressures[earliest.segmentID]
		if m.reachBlockingL0Threshold(pressure.rows, pressure.bytes) {
			sealSegmentIDs[earliest.segmentID] = policy.PolicyBlockingL0(pressure.rows, pressure.bytes, m.cfg.blockingL0EntryNum, m.cfg.blockingL0SizeBytes)
		}
	}
	if len(sealSegmentIDs) == 0 {
		return nil
	}
	return sealSegmentIDs
}

func (m *StatsManager) reachBlockingL0Threshold(rows uint64, bytes uint64) bool {
	if m.cfg.blockingL0EntryNum >= 0 && rows >= uint64(m.cfg.blockingL0EntryNum) {
		return true
	}
	if m.cfg.blockingL0SizeBytes >= 0 && bytes >= uint64(m.cfg.blockingL0SizeBytes) {
		return true
	}
	return false
}

// selectSegmentsUntilLessThanLWM selects segments until the total size is less than the threshold.
func (m *StatsManager) selectSegmentsUntilLessThanLWM() []int64 {
	m.mu.Lock()
	restSpace := utils.SaturatingUint64ToInt64(m.totalFlushSize) - m.cfg.growingBytesLWM
	m.mu.Unlock()

	if restSpace <= 0 {
		return nil
	}

	segmentIDs := make([]int64, 0)
	stats := m.createStatsSlice()
	statsHeap := typeutil.NewObjectArrayBasedMaximumHeap(stats, func(s segmentWithBinarySize) uint64 {
		return s.binarySize
	})
	for restSpace > 0 && statsHeap.Len() > 0 {
		nextOne := statsHeap.Pop()
		restSpace -= int64(nextOne.binarySize)
		segmentIDs = append(segmentIDs, nextOne.segmentID)
	}
	return segmentIDs
}

// createStatsSlice creates a slice of SegmentStats.
func (m *StatsManager) createStatsSlice() []segmentWithBinarySize {
	m.mu.Lock()
	defer m.mu.Unlock()

	stats := make([]segmentWithBinarySize, 0, len(m.segmentStats))
	for id, stat := range m.segmentStats {
		flushSize := stat.FlushSize()
		if flushSize > 0 {
			stats = append(stats, segmentWithBinarySize{
				segmentID:  id,
				binarySize: flushSize,
			})
		}
	}
	return stats
}

// updateConfig updates the config of stats manager.
func (m *StatsManager) updateConfig() {
	cfg := newStatsConfig()
	if err := cfg.Validate(); err != nil {
		return
	}
	m.metricHelper.ObserveConfigUpdate(cfg)

	m.mu.Lock()
	defer m.mu.Unlock()
	if m.cfg != cfg {
		m.Logger().Info(context.TODO(), "update stats manager config", mlog.Any("newConfig", cfg), mlog.Any("oldConfig", m.cfg))
		m.cfg = cfg
	}
}

// getConfig gets the config of stats manager.
func (m *StatsManager) getConfig() statsConfig {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.cfg
}

type segmentWithBinarySize struct {
	segmentID  int64
	binarySize uint64
}
