// 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 measure

import (
	"context"
	"fmt"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/pkg/errors"
	"go.uber.org/multierr"

	"github.com/apache/skywalking-banyandb/api/common"
	"github.com/apache/skywalking-banyandb/api/data"
	commonv1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/common/v1"
	databasev1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/database/v1"
	measurev1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/measure/v1"
	"github.com/apache/skywalking-banyandb/banyand/internal/storage"
	"github.com/apache/skywalking-banyandb/banyand/metadata"
	"github.com/apache/skywalking-banyandb/banyand/metadata/schema"
	"github.com/apache/skywalking-banyandb/banyand/observability"
	obsservice "github.com/apache/skywalking-banyandb/banyand/observability/services"
	"github.com/apache/skywalking-banyandb/banyand/protector"
	"github.com/apache/skywalking-banyandb/banyand/queue"
	"github.com/apache/skywalking-banyandb/pkg/bus"
	"github.com/apache/skywalking-banyandb/pkg/fs"
	"github.com/apache/skywalking-banyandb/pkg/logger"
	"github.com/apache/skywalking-banyandb/pkg/meter/native"
	banyandbpath "github.com/apache/skywalking-banyandb/pkg/path"
	"github.com/apache/skywalking-banyandb/pkg/run"
	resourceSchema "github.com/apache/skywalking-banyandb/pkg/schema"
	"github.com/apache/skywalking-banyandb/pkg/timestamp"
)

const (
	serviceName = "measure"
)

var _ Service = (*dataSVC)(nil)

type dataSVC struct {
	lfs                fs.FileSystem
	c                  storage.Cache
	pipeline           queue.Server
	omr                observability.MetricsRegistry
	metadata           metadata.Repo
	pm                 protector.Memory
	metricPipeline     queue.Server
	l                  *logger.Logger
	schemaRepo         *schemaRepo
	cm                 *cacheMetrics
	diskMonitor        *storage.DiskMonitor
	root               string
	dataPath           string
	snapshotDir        string
	option             option
	retentionConfig    storage.RetentionConfig
	cc                 storage.CacheConfig
	maxFileSnapshotNum int
	minFileSnapshotAge time.Duration
}

func (s *dataSVC) Measure(metadata *commonv1.Metadata) (Measure, error) {
	sm, ok := s.schemaRepo.loadMeasure(metadata)
	if !ok {
		return nil, errors.WithStack(ErrMeasureNotExist)
	}
	return sm, nil
}

func (s *dataSVC) LoadGroup(name string) (resourceSchema.Group, bool) {
	return s.schemaRepo.LoadGroup(name)
}

func (s *dataSVC) DropGroup(_ context.Context, groupName string) error {
	return s.schemaRepo.DropGroup(groupName)
}

func (s *dataSVC) GetRemovalSegmentsTimeRange(group string) *timestamp.TimeRange {
	return s.schemaRepo.GetRemovalSegmentsTimeRange(group)
}

// RetentionService interface implementation.

func (s *dataSVC) GetDataPath() string {
	return s.dataPath
}

func (s *dataSVC) GetSnapshotDir() string {
	return s.snapshotDir
}

func (s *dataSVC) LoadAllGroups() []resourceSchema.Group {
	return s.schemaRepo.LoadAllGroups()
}

func (s *dataSVC) PeekOldestSegmentEndTimeInGroup(group string) (time.Time, bool) {
	g, ok := s.schemaRepo.LoadGroup(group)
	if !ok {
		return time.Time{}, false
	}

	db := g.SupplyTSDB()
	if db == nil {
		return time.Time{}, false
	}

	// Type assert to the storage interface that has PeekOldestSegmentEndTime
	if dbWithPeek, ok := db.(interface{ PeekOldestSegmentEndTime() (time.Time, bool) }); ok {
		return dbWithPeek.PeekOldestSegmentEndTime()
	}

	return time.Time{}, false
}

func (s *dataSVC) DeleteOldestSegmentInGroup(group string) (bool, error) {
	g, ok := s.schemaRepo.LoadGroup(group)
	if !ok {
		return false, nil
	}

	db := g.SupplyTSDB()
	if db == nil {
		return false, nil
	}

	// Type assert to the storage interface that has DeleteOldestSegment
	if dbWithDelete, ok := db.(interface{ DeleteOldestSegment() (bool, error) }); ok {
		return dbWithDelete.DeleteOldestSegment()
	}

	s.l.Debug().Str("group", group).Msg("database does not support DeleteOldestSegment")
	return false, nil
}

func (s *dataSVC) CleanupOldSnapshots(maxAge time.Duration) error {
	if s.snapshotDir == "" {
		return nil
	}
	// Use age-based cleanup during forced retention cleanup
	storage.DeleteOldSnapshots(s.snapshotDir, maxAge, s.lfs)
	return nil
}

func (s *dataSVC) GetServiceName() string {
	return serviceName
}

func (s *dataSVC) FlagSet() *run.FlagSet {
	flagS := run.NewFlagSet("storage")
	flagS.StringVar(&s.root, "measure-root-path", "/tmp", "the root path of measure")
	flagS.StringVar(&s.dataPath, "measure-data-path", "", "the data directory path of measure. If not set, <measure-root-path>/measure/data will be used")
	flagS.DurationVar(&s.option.flushTimeout, "measure-flush-timeout", defaultFlushTimeout, "the memory data timeout of measure")
	s.option.mergePolicy = newDefaultMergePolicy()
	flagS.VarP(&s.option.mergePolicy.maxFanOutSize, "measure-max-fan-out-size", "", "the upper bound of a single file size after merge of measure")
	s.option.seriesCacheMaxSize = run.Bytes(32 << 20)
	flagS.VarP(&s.option.seriesCacheMaxSize, "measure-series-cache-max-size", "", "the max size of series cache in each group")

	// Retention configuration flags
	flagS.Float64Var(&s.retentionConfig.HighWatermark, "measure-retention-high-watermark", 95.0, "disk usage high watermark for forced retention cleanup")
	flagS.Float64Var(&s.retentionConfig.LowWatermark, "measure-retention-low-watermark", 85.0, "disk usage low watermark percentage where forced retention cleanup stops")
	flagS.DurationVar(&s.retentionConfig.CheckInterval, "measure-retention-check-interval", 5*time.Minute, "interval for checking disk usage")
	flagS.DurationVar(&s.retentionConfig.Cooldown, "measure-retention-cooldown", 30*time.Second, "cooldown period between forced segment deletions")
	flagS.BoolVar(&s.retentionConfig.ForceCleanupEnabled, "measure-retention-force-cleanup-enabled", false,
		"enable forced retention cleanup when disk usage exceeds high watermark")

	flagS.IntVar(&s.maxFileSnapshotNum, "measure-max-file-snapshot-num", 10, "the maximum number of file snapshots allowed")
	flagS.DurationVar(&s.minFileSnapshotAge, "measure-min-file-snapshot-age", time.Hour, "minimum age for file snapshots to be eligible for deletion")
	s.cc.MaxCacheSize = run.Bytes(100 * 1024 * 1024)
	flagS.VarP(&s.cc.MaxCacheSize, "service-cache-max-size", "", "maximum service cache size (e.g., 100M)")
	flagS.DurationVar(&s.cc.CleanupInterval, "service-cache-cleanup-interval", 30*time.Second, "service cache cleanup interval")
	flagS.DurationVar(&s.cc.IdleTimeout, "service-cache-idle-timeout", 2*time.Minute, "service cache entry idle timeout")
	return flagS
}

func (s *dataSVC) Validate() error {
	if s.root == "" {
		return errEmptyRootPath
	}

	// Validate retention configuration
	if s.retentionConfig.HighWatermark < 0 || s.retentionConfig.HighWatermark > 100 {
		return errors.New("measure-retention-high-watermark must be between 0 and 100")
	}
	if s.retentionConfig.LowWatermark < 0 || s.retentionConfig.LowWatermark > 100 {
		return errors.New("measure-retention-low-watermark must be between 0 and 100")
	}
	if s.retentionConfig.LowWatermark > s.retentionConfig.HighWatermark {
		return errors.New("measure-retention-low-watermark must be less than measure-retention-high-watermark")
	}
	if s.retentionConfig.CheckInterval <= 0 && s.retentionConfig.ForceCleanupEnabled {
		return errors.New("measure-retention-check-interval must be greater than 0 when force cleanup is enabled")
	}
	if s.retentionConfig.Cooldown <= 0 {
		return errors.New("measure-retention-cooldown must be greater than 0")
	}

	if s.cc.MaxCacheSize < 0 {
		return errors.New("service-cache-max-size must be greater than or equal to 0")
	}
	if s.cc.CleanupInterval <= 0 {
		return errors.New("service-cache-cleanup-interval must be greater than 0")
	}
	if s.cc.IdleTimeout <= 0 {
		return errors.New("service-cache-idle-timeout must be greater than 0")
	}
	return nil
}

func (s *dataSVC) Name() string {
	return serviceName
}

func (s *dataSVC) Role() databasev1.Role {
	return databasev1.Role_ROLE_DATA
}

func (s *dataSVC) PreRun(ctx context.Context) error {
	s.l = logger.GetLogger(s.Name())
	s.l.Info().Msg("memory protector is initialized in PreRun")
	s.lfs = fs.NewLocalFileSystemWithLoggerAndLimit(s.l, s.pm.GetLimit())
	var err error
	if s.root, err = banyandbpath.Get(s.root); err != nil {
		return err
	}
	if s.dataPath != "" {
		if s.dataPath, err = banyandbpath.Get(s.dataPath); err != nil {
			return err
		}
	}
	path := path.Join(s.root, s.Name())
	s.snapshotDir = filepath.Join(path, storage.SnapshotsDir)
	obsservice.UpdatePath(path)
	if s.dataPath == "" {
		s.dataPath = filepath.Join(path, storage.DataDir)
	}
	if !strings.HasPrefix(filepath.VolumeName(s.dataPath), filepath.VolumeName(path)) {
		obsservice.UpdatePath(s.dataPath)
	}
	val := ctx.Value(common.ContextNodeKey)
	if val == nil {
		return errors.New("node id is empty")
	}
	if s.cc.MaxCacheSize == 0 {
		s.c = storage.NewBypassCache()
	} else {
		s.c = storage.NewServiceCacheWithConfig(s.cc)
	}
	node := val.(common.Node)
	s.schemaRepo = newDataSchemaRepo(s.dataPath, s, node.Labels, node.NodeID)

	s.cm = newCacheMetrics(s.omr)
	obsservice.MetricsCollector.Register("measure_cache", s.collectCacheMetrics)

	if s.pipeline == nil {
		return nil
	}

	collectDataInfoListener := &collectDataInfoListener{s: s}
	if subscribeErr := s.pipeline.Subscribe(data.TopicMeasureCollectDataInfo, collectDataInfoListener); subscribeErr != nil {
		return fmt.Errorf("failed to subscribe to collect data info topic: %w", subscribeErr)
	}
	if dropGroupErr := s.pipeline.Subscribe(data.TopicMeasureDropGroup, &dropGroupDataListener{s: s}); dropGroupErr != nil {
		return fmt.Errorf("failed to subscribe to drop group topic: %w", dropGroupErr)
	}

	if err = s.createDataNativeObservabilityGroup(ctx); err != nil {
		return err
	}

	if err = s.pipeline.Subscribe(data.TopicSnapshot, &dataSnapshotListener{s: s}); err != nil {
		return err
	}

	if err = s.pipeline.Subscribe(data.TopicMeasureDeleteExpiredSegments, &dataDeleteStreamSegmentsListener{s: s}); err != nil {
		return err
	}

	s.pipeline.RegisterChunkedSyncHandler(data.TopicMeasurePartSync, setUpChunkedSyncCallback(s.l, s.schemaRepo))
	s.pipeline.RegisterChunkedSyncHandler(data.TopicMeasureSeriesSync, setUpSyncSeriesCallback(s.l, s.schemaRepo))
	err = s.pipeline.Subscribe(data.TopicMeasureSeriesIndexInsert, setUpIndexCallback(s.l, s.schemaRepo, data.TopicMeasureSeriesIndexInsert))
	if err != nil {
		return err
	}
	err = s.pipeline.Subscribe(data.TopicMeasureSeriesIndexUpdate, setUpIndexCallback(s.l, s.schemaRepo, data.TopicMeasureSeriesIndexUpdate))
	if err != nil {
		return err
	}

	// Initialize disk monitor for forced retention
	s.diskMonitor = storage.NewDiskMonitor(s, s.retentionConfig, s.omr)
	s.diskMonitor.Start()

	// For now, keep the original write throttling behavior based on high watermark
	// TODO: Replace this with a newer write callback that doesn't duplicate disk monitoring
	writeListener := setUpWriteCallback(s.l, s.schemaRepo, int(s.retentionConfig.HighWatermark))
	err = s.pipeline.Subscribe(data.TopicMeasureWrite, writeListener)
	if err != nil {
		return err
	}
	// only subscribe metricPipeline for data node
	if s.metricPipeline != nil {
		err = s.metricPipeline.Subscribe(data.TopicMeasureWrite, writeListener)
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *dataSVC) Serve() run.StopNotify {
	return s.schemaRepo.StopCh()
}

func (s *dataSVC) GracefulStop() {
	// Stop disk monitor
	if s.diskMonitor != nil {
		s.diskMonitor.Stop()
	}

	obsservice.MetricsCollector.Unregister("measure_cache")
	s.schemaRepo.Close()
	s.c.Close()
}

func (s *dataSVC) collectCacheMetrics() {
	if s.cm == nil || s.c == nil {
		return
	}

	requests := s.c.Requests()
	misses := s.c.Misses()
	length := s.c.Entries()
	size := s.c.Size()
	var hitRatio float64
	if requests > 0 {
		hitRatio = float64(requests-misses) / float64(requests)
	}

	s.cm.requests.Set(float64(requests))
	s.cm.misses.Set(float64(misses))
	s.cm.hitRatio.Set(hitRatio)
	s.cm.entries.Set(float64(length))
	s.cm.size.Set(float64(size))
}

func (s *dataSVC) createDataNativeObservabilityGroup(ctx context.Context) error {
	if !s.omr.NativeEnabled() {
		return nil
	}
	g := &commonv1.Group{
		Metadata: &commonv1.Metadata{
			Name: native.ObservabilityGroupName,
		},
		Catalog: commonv1.Catalog_CATALOG_MEASURE,
		ResourceOpts: &commonv1.ResourceOpts{
			ShardNum: 1,
			SegmentInterval: &commonv1.IntervalRule{
				Unit: commonv1.IntervalRule_UNIT_DAY,
				Num:  1,
			},
			Ttl: &commonv1.IntervalRule{
				Unit: commonv1.IntervalRule_UNIT_DAY,
				Num:  1,
			},
		},
	}
	if err := s.metadata.GroupRegistry().CreateGroup(ctx, g); err != nil &&
		!errors.Is(err, schema.ErrGRPCAlreadyExists) {
		return errors.WithMessage(err, "failed to create native observability group")
	}
	return nil
}

func (s *dataSVC) takeGroupSnapshot(dstDir string, groupName string) (bool, error) {
	group, ok := s.schemaRepo.LoadGroup(groupName)
	if !ok {
		return false, errors.Errorf("group %s not found", groupName)
	}
	db := group.SupplyTSDB()
	if db == nil {
		return false, errors.Errorf("group %s has no tsdb", group.GetSchema().Metadata.Name)
	}
	tsdb := db.(storage.TSDB[*tsTable, option])
	created, err := tsdb.TakeFileSnapshot(dstDir)
	if err != nil {
		return false, errors.WithMessagef(err, "snapshot %s fail to take file snapshot for group %s", dstDir, group.GetSchema().Metadata.Name)
	}
	return created, nil
}

// NewDataSVC returns a new data service.
func NewDataSVC(metadata metadata.Repo, pipeline queue.Server, metricPipeline queue.Server, omr observability.MetricsRegistry, pm protector.Memory) (Service, error) {
	return &dataSVC{
		metadata:       metadata,
		pipeline:       pipeline,
		metricPipeline: metricPipeline,
		omr:            omr,
		pm:             pm,
	}, nil
}

// NewReadonlyDataSVC returns a new readonly data service.
func NewReadonlyDataSVC(metadata metadata.Repo, omr observability.MetricsRegistry, pm protector.Memory) (Service, error) {
	return &dataSVC{
		metadata: metadata,
		omr:      omr,
		pm:       pm,
	}, nil
}

func newDataSchemaRepo(path string, svc *dataSVC, nodeLabels map[string]string, nodeID string) *schemaRepo {
	sr := &schemaRepo{
		path:     path,
		l:        svc.l,
		metadata: svc.metadata,
		nodeID:   nodeID,
		role:     databasev1.Role_ROLE_DATA,
	}
	sr.Repository = resourceSchema.NewRepository(
		svc.metadata,
		svc.l,
		newDataSupplier(path, svc, sr, nodeLabels),
		resourceSchema.NewMetrics(svc.omr.With(metadataScope)),
	)
	sr.start()
	return sr
}

func newDataSupplier(path string, svc *dataSVC, sr *schemaRepo, nodeLabels map[string]string) *supplier {
	if svc.pm == nil {
		svc.l.Panic().Msg("CRITICAL: svc.pm is nil in newSupplier")
	}
	opt := svc.option
	opt.protector = svc.pm

	if opt.protector == nil {
		svc.l.Panic().Msg("CRITICAL: opt.protector is still nil after assignment")
	}
	return &supplier{
		path:       path,
		metadata:   svc.metadata,
		l:          svc.l,
		c:          svc.c,
		option:     opt,
		omr:        svc.omr,
		pm:         svc.pm,
		schemaRepo: sr,
		nodeLabels: nodeLabels,
	}
}

type dataSnapshotListener struct {
	*bus.UnImplementedHealthyListener
	s           *dataSVC
	snapshotSeq uint64
	snapshotMux sync.Mutex
}

func (d *dataSnapshotListener) Rev(ctx context.Context, message bus.Message) bus.Message {
	groups := message.Data().([]*databasev1.SnapshotRequest_Group)
	var gg []resourceSchema.Group
	if len(groups) == 0 {
		gg = d.s.schemaRepo.LoadAllGroups()
		n := ""
		for _, g := range gg {
			if n != "" {
				n += ","
			}
			n += g.GetSchema().Metadata.Name
		}

		log.Info().Msgf("loaded all snapshots: %s", n)
	} else {
		for _, g := range groups {
			if g.Catalog != commonv1.Catalog_CATALOG_MEASURE {
				continue
			}
			group, ok := d.s.schemaRepo.LoadGroup(g.Group)
			if !ok {
				continue
			}
			gg = append(gg, group)
		}
		log.Info().Msgf("loaded groups: %s", gg)
	}
	if len(gg) == 0 {
		return bus.NewMessage(bus.MessageID(time.Now().UnixNano()), nil)
	}
	d.snapshotMux.Lock()
	defer d.snapshotMux.Unlock()
	storage.DeleteStaleSnapshots(d.s.snapshotDir, d.s.maxFileSnapshotNum, d.s.minFileSnapshotAge, d.s.lfs)
	sn := d.snapshotName()
	var err error
	var snapshotCreated int
	for _, g := range gg {
		select {
		case <-ctx.Done():
			return bus.NewMessage(bus.MessageID(time.Now().UnixNano()), nil)
		default:
		}
		groupName := g.GetSchema().Metadata.Name
		snapshotPath := filepath.Join(d.s.snapshotDir, sn, groupName)
		created, errGroup := d.s.takeGroupSnapshot(snapshotPath, groupName)
		if errGroup != nil {
			d.s.l.Error().Err(errGroup).Str("group", groupName).Msg("fail to take group snapshot")
			err = multierr.Append(err, errGroup)
			continue
		}
		if !created {
			d.s.l.Info().Str("group", groupName).Msg("skip empty group snapshot")
			continue
		}
		snapshotCreated++
	}
	if snapshotCreated == 0 && err == nil {
		return bus.NewMessage(bus.MessageID(time.Now().UnixNano()), nil)
	}
	snp := &databasev1.Snapshot{
		Name:    sn,
		Catalog: commonv1.Catalog_CATALOG_MEASURE,
	}
	log.Info().Msgf("snapshot %s created", sn)
	if err != nil {
		snp.Error = err.Error()
	}
	return bus.NewMessage(bus.MessageID(time.Now().UnixNano()), snp)
}

func (d *dataSnapshotListener) snapshotName() string {
	d.snapshotSeq++
	return fmt.Sprintf("%s-%08X", time.Now().UTC().Format(storage.SnapshotTimeFormat), d.snapshotSeq)
}

type dataDeleteStreamSegmentsListener struct {
	*bus.UnImplementedHealthyListener
	s *dataSVC
}

func (d *dataDeleteStreamSegmentsListener) Rev(_ context.Context, message bus.Message) bus.Message {
	req := message.Data().(*measurev1.DeleteExpiredSegmentsRequest)
	if req == nil {
		return bus.NewMessage(bus.MessageID(time.Now().UnixNano()), int64(0))
	}

	db, err := d.s.schemaRepo.loadTSDB(req.Group)
	if err != nil {
		d.s.l.Error().Err(err).Str("group", req.Group).Msg("failed to load tsdb")
		return bus.NewMessage(bus.MessageID(time.Now().UnixNano()), int64(0))
	}
	deleted := db.DeleteExpiredSegments(req.SegmentSuffixes)
	return bus.NewMessage(bus.MessageID(time.Now().UnixNano()), deleted)
}

func (s *dataSVC) CollectDataInfo(ctx context.Context, group string) (*databasev1.DataInfo, error) {
	return s.schemaRepo.CollectDataInfo(ctx, group)
}

func (s *dataSVC) CollectLiaisonInfo(_ context.Context, _ string) (*databasev1.LiaisonInfo, error) {
	return nil, errors.New("collect liaison info is not supported on data node")
}

type collectDataInfoListener struct {
	*bus.UnImplementedHealthyListener
	s *dataSVC
}

func (l *collectDataInfoListener) Rev(ctx context.Context, message bus.Message) bus.Message {
	req, ok := message.Data().(*databasev1.GroupRegistryServiceInspectRequest)
	if !ok {
		return bus.NewMessage(message.ID(), common.NewError("invalid data type for collect data info request"))
	}
	dataInfo, collectErr := l.s.schemaRepo.CollectDataInfo(ctx, req.Group)
	if collectErr != nil {
		return bus.NewMessage(message.ID(), common.NewError("failed to collect data info: %v", collectErr))
	}
	return bus.NewMessage(message.ID(), dataInfo)
}

type dropGroupDataListener struct {
	*bus.UnImplementedHealthyListener
	s *dataSVC
}

func (l *dropGroupDataListener) Rev(ctx context.Context, message bus.Message) bus.Message {
	req, ok := message.Data().(*databasev1.GroupRegistryServiceDeleteRequest)
	if !ok {
		return bus.NewMessage(message.ID(), common.NewError("invalid data type for drop group request"))
	}
	if dropErr := l.s.DropGroup(ctx, req.Group); dropErr != nil {
		return bus.NewMessage(message.ID(), common.NewError("failed to drop group data: %v", dropErr))
	}
	return bus.NewMessage(message.ID(), req)
}
