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

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort" // added for sorting parts
	"strings"
	"sync"
	"sync/atomic"
	"time"

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

	commonv1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/common/v1"
	databasev1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/database/v1"
	"github.com/apache/skywalking-banyandb/banyand/internal/storage"
	"github.com/apache/skywalking-banyandb/pkg/bus"
	"github.com/apache/skywalking-banyandb/pkg/logger"
	"github.com/apache/skywalking-banyandb/pkg/schema"
)

func (tst *tsTable) currentSnapshot() *snapshot {
	tst.RLock()
	defer tst.RUnlock()
	if tst.snapshot == nil {
		return nil
	}
	s := tst.snapshot
	s.incRef()
	return s
}

type snapshotCreator rune

const (
	snapshotCreatorMemPart = iota
	snapshotCreatorFlusher
	snapshotCreatorMerger
	snapshotCreatorMergedFlusher
	snapshotCreatorSyncer
)

type snapshot struct {
	parts   []*partWrapper
	epoch   uint64
	creator snapshotCreator

	ref int32
}

func (s *snapshot) getParts(dst []*part, minTimestamp, maxTimestamp int64) ([]*part, int) {
	var count int
	for _, p := range s.parts {
		pm := p.p.partMetadata
		if maxTimestamp < pm.MinTimestamp || minTimestamp > pm.MaxTimestamp {
			continue
		}
		dst = append(dst, p.p)
		count++
	}
	return dst, count
}

func (s *snapshot) incRef() {
	atomic.AddInt32(&s.ref, 1)
}

func (s *snapshot) decRef() {
	n := atomic.AddInt32(&s.ref, -1)
	if n > 0 {
		return
	}
	for i := range s.parts {
		s.parts[i].decRef()
	}
	s.parts = s.parts[:0]
}

func (s *snapshot) copyAllTo(nextEpoch uint64) snapshot {
	var result snapshot
	result.epoch = nextEpoch
	result.ref = 1
	for i := range s.parts {
		s.parts[i].incRef()
		result.parts = append(result.parts, s.parts[i])
	}
	return result
}

func (s *snapshot) merge(nextEpoch uint64, nextParts map[uint64]*partWrapper) snapshot {
	var result snapshot
	result.epoch = nextEpoch
	result.ref = 1
	for i := 0; i < len(s.parts); i++ {
		if n, ok := nextParts[s.parts[i].ID()]; ok {
			result.parts = append(result.parts, n)
			continue
		}
		s.parts[i].incRef()
		result.parts = append(result.parts, s.parts[i])
	}
	return result
}

func (s *snapshot) remove(nextEpoch uint64, merged map[uint64]struct{}) snapshot {
	var result snapshot
	result.epoch = nextEpoch
	result.ref = 1
	for i := 0; i < len(s.parts); i++ {
		if _, ok := merged[s.parts[i].ID()]; !ok {
			s.parts[i].incRef()
			result.parts = append(result.parts, s.parts[i])
			continue
		}
		s.parts[i].removable.Store(true)
	}
	return result
}

func getDisjointParts(parts []*part, asc bool) [][]*part {
	if len(parts) == 0 {
		return nil
	}
	sort.Slice(parts, func(i, j int) bool {
		return parts[i].partMetadata.MinTimestamp < parts[j].partMetadata.MinTimestamp
	})

	var groups [][]*part
	var currentGroup []*part
	var boundary int64
	for _, p := range parts {
		pMin := p.partMetadata.MinTimestamp
		pMax := p.partMetadata.MaxTimestamp
		if len(currentGroup) == 0 {
			currentGroup = append(currentGroup, p)
			boundary = pMax
		} else {
			if pMin <= boundary {
				currentGroup = append(currentGroup, p)
				if pMax > boundary {
					boundary = pMax
				}
			} else {
				groups = append(groups, currentGroup)
				currentGroup = []*part{p}
				boundary = pMax
			}
		}
	}

	if len(currentGroup) > 0 {
		groups = append(groups, currentGroup)
	}

	if !asc {
		for i, j := 0, len(groups)-1; i < j; i, j = i+1, j-1 {
			groups[i], groups[j] = groups[j], groups[i]
		}
	}
	return groups
}

func snapshotName(snapshot uint64) string {
	return fmt.Sprintf("%016x%s", snapshot, snapshotSuffix)
}

func parseSnapshot(name string) (uint64, error) {
	if filepath.Ext(name) != snapshotSuffix {
		return 0, errors.New("invalid snapshot file ext")
	}
	if len(name) < 16 {
		return 0, errors.New("invalid snapshot file name")
	}
	return parseEpoch(name[:16])
}

func (tst *tsTable) TakeFileSnapshot(dst string) (success bool, err error) {
	if tst.index == nil {
		return false, fmt.Errorf("cannot take file snapshot: index is not initialized for this tsTable")
	}
	snapshot := tst.currentSnapshot()
	if snapshot == nil {
		return false, storage.ErrNoCurrentSnapshot
	}
	defer snapshot.decRef()
	defer func() {
		if err != nil {
			tst.fileSystem.MustRMAll(dst)
		}
	}()
	indexDir := filepath.Join(dst, filepath.Base(tst.index.location))
	tst.fileSystem.MkdirPanicIfExist(indexDir, storage.DirPerm)
	if indexErr := tst.index.store.TakeFileSnapshot(indexDir); indexErr != nil {
		return false, fmt.Errorf("failed to take file snapshot for index: %w", indexErr)
	}

	hasDiskParts := false
	for _, pw := range snapshot.parts {
		if pw.mp != nil {
			continue
		}

		part := pw.p
		srcPath := part.path
		destPartPath := filepath.Join(dst, filepath.Base(srcPath))

		if linkErr := tst.fileSystem.CreateHardLink(srcPath, destPartPath, nil); linkErr != nil {
			return false, fmt.Errorf("failed to create snapshot for part %d: %w", part.partMetadata.ID, linkErr)
		}
		hasDiskParts = true
	}
	if !hasDiskParts {
		return true, nil
	}
	tst.createMetadata(dst, snapshot)
	parent := filepath.Dir(dst)
	tst.fileSystem.SyncPath(parent)
	return true, nil
}

func (tst *tsTable) createMetadata(dst string, snapshot *snapshot) {
	var partNames []string
	for i := range snapshot.parts {
		partNames = append(partNames, partName(snapshot.parts[i].ID()))
	}
	data, err := json.Marshal(partNames)
	if err != nil {
		logger.Panicf("cannot marshal partNames to JSON: %s", err)
	}
	snapshotPath := filepath.Join(dst, snapshotName(snapshot.epoch))
	lf, err := tst.fileSystem.CreateFile(snapshotPath, storage.FilePerm)
	if err != nil {
		logger.Panicf("cannot create lock file %s: %s", snapshotPath, err)
	}
	n, err := lf.Write(data)
	if err != nil {
		logger.Panicf("cannot write snapshot %s: %s", snapshotPath, err)
	}
	if n != len(data) {
		logger.Panicf("unexpected number of bytes written to %s; got %d; want %d", snapshotPath, n, len(data))
	}
}

func (s *standalone) 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
}

// collectSegDirs walks a directory tree and collects all seg-* directory paths.
// It only collects directories matching the "seg-*" pattern and ignores all files.
// Returns a map of relative paths to seg-* directories.
func collectSegDirs(rootDir string) (map[string]bool, error) {
	segDirs := make(map[string]bool)

	walkErr := filepath.Walk(rootDir, func(currentPath string, info os.FileInfo, err error) error {
		if err != nil {
			// If we can't read a directory, log but continue
			return nil
		}

		// Skip files, only process directories
		if !info.IsDir() {
			return nil
		}

		// Get the directory name
		dirName := filepath.Base(currentPath)

		// Check if this is a seg-* directory
		if strings.HasPrefix(dirName, "seg-") {
			// Get relative path from root
			relPath, relErr := filepath.Rel(rootDir, currentPath)
			if relErr != nil {
				return fmt.Errorf("failed to get relative path for %s: %w", currentPath, relErr)
			}

			// Add to our collection
			segDirs[relPath] = true

			// Don't recurse into seg-* directories
			return filepath.SkipDir
		}

		return nil
	})

	if walkErr != nil {
		return nil, fmt.Errorf("failed to walk directory %s: %w", rootDir, walkErr)
	}

	return segDirs, nil
}

// compareSegDirs compares two sets of seg-* directories and returns three slices:
// matched (in both), onlyInSnapshot (only in snapshot), onlyInData (only in data).
func compareSegDirs(snapshotDirs, dataDirs map[string]bool) (matched, onlyInSnapshot, onlyInData []string) {
	// Find directories in both sets (matched)
	for dir := range snapshotDirs {
		if dataDirs[dir] {
			matched = append(matched, dir)
		} else {
			onlyInSnapshot = append(onlyInSnapshot, dir)
		}
	}

	// Find directories only in data
	for dir := range dataDirs {
		if !snapshotDirs[dir] {
			onlyInData = append(onlyInData, dir)
		}
	}

	// Sort for consistent output
	sort.Strings(matched)
	sort.Strings(onlyInSnapshot)
	sort.Strings(onlyInData)

	return matched, onlyInSnapshot, onlyInData
}

// compareSnapshotWithData compares the snapshot directory with the data directory
// to verify that seg-* directories are consistent. Only logs differences, does not return errors.
func (s *snapshotListener) compareSnapshotWithData(snapshotDir, dataDir, groupName string) {
	// Collect seg-* directories from snapshot
	snapshotDirs, snapshotErr := collectSegDirs(snapshotDir)
	if snapshotErr != nil {
		s.s.l.Warn().Err(snapshotErr).
			Str("group", groupName).
			Str("snapshotDir", snapshotDir).
			Msg("failed to collect seg-* directories from snapshot, skipping comparison")
		return
	}

	// Collect seg-* directories from data
	dataDirs, dataErr := collectSegDirs(dataDir)
	if dataErr != nil {
		s.s.l.Warn().Err(dataErr).
			Str("group", groupName).
			Str("dataDir", dataDir).
			Msg("failed to collect seg-* directories from data, skipping comparison")
		return
	}

	// Compare the directories
	matched, onlyInSnapshot, onlyInData := compareSegDirs(snapshotDirs, dataDirs)

	// Log consolidated comparison results at Info level
	s.s.l.Info().
		Str("group", groupName).
		Int("matched", len(matched)).
		Int("onlyInSnapshot", len(onlyInSnapshot)).
		Int("onlyInData", len(onlyInData)).
		Strs("matchedDirs", matched).
		Strs("onlyInSnapshotDirs", onlyInSnapshot).
		Strs("onlyInDataDirs", onlyInData).
		Msgf("snapshot comparison for group %s, data dir: %s, snapshot dir: %s",
			groupName, dataDir, snapshotDir)
}

type snapshotListener struct {
	*bus.UnImplementedHealthyListener
	s           *standalone
	snapshotSeq uint64
	snapshotMux sync.Mutex
}

// Rev takes a snapshot of the database.
func (s *snapshotListener) Rev(ctx context.Context, message bus.Message) bus.Message {
	groups := message.Data().([]*databasev1.SnapshotRequest_Group)
	var gg []schema.Group
	if len(groups) == 0 {
		gg = s.s.schemaRepo.LoadAllGroups()
	} else {
		for _, g := range groups {
			if g.Catalog != commonv1.Catalog_CATALOG_STREAM {
				continue
			}
			group, ok := s.s.schemaRepo.LoadGroup(g.Group)
			if !ok {
				continue
			}
			gg = append(gg, group)
		}
	}
	if len(gg) == 0 {
		return bus.NewMessage(bus.MessageID(time.Now().UnixNano()), nil)
	}
	s.snapshotMux.Lock()
	defer s.snapshotMux.Unlock()
	storage.DeleteStaleSnapshots(s.s.snapshotDir, s.s.maxFileSnapshotNum, s.s.minFileSnapshotAge, s.s.lfs)
	sn := s.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(s.s.snapshotDir, sn, groupName)
		created, errGroup := s.s.takeGroupSnapshot(snapshotPath, groupName)
		if errGroup != nil {
			s.s.l.Error().Err(errGroup).Str("group", groupName).Msg("fail to take group snapshot")
			err = multierr.Append(err, errGroup)
			continue
		}
		if !created {
			s.s.l.Info().Str("group", groupName).Msg("skip empty group snapshot")
			continue
		}
		snapshotCreated++

		// Compare snapshot with data directory to verify consistency
		dataPath := filepath.Join(s.s.dataPath, groupName)
		s.compareSnapshotWithData(snapshotPath, dataPath, groupName)
	}
	if snapshotCreated == 0 && err == nil {
		return bus.NewMessage(bus.MessageID(time.Now().UnixNano()), nil)
	}
	snp := &databasev1.Snapshot{
		Name:    sn,
		Catalog: commonv1.Catalog_CATALOG_STREAM,
	}
	if err != nil {
		snp.Error = err.Error()
	}
	return bus.NewMessage(bus.MessageID(time.Now().UnixNano()), snp)
}

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