// Copyright 2019 Dolthub, 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 pull

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

	"golang.org/x/sync/errgroup"

	"github.com/dolthub/dolt/go/libraries/doltcore/dconfig"
	"github.com/dolthub/dolt/go/store/chunks"
	"github.com/dolthub/dolt/go/store/hash"
	"github.com/dolthub/dolt/go/store/nbs"
)

// ErrDBUpToDate is the error code returned from NewPuller in the event that there is no work to do.
var ErrDBUpToDate = errors.New("the database does not need to be pulled as it's already up to date")

// ErrIncompatibleSourceChunkStore is the error code returned from NewPuller in
// the event that the source ChunkStore does not implement `NBSCompressedChunkStore`.
var ErrIncompatibleSourceChunkStore = errors.New("the chunk store of the source database does not implement NBSCompressedChunkStore.")

type WalkAddrs func(chunks.Chunk, func(hash.Hash, bool) error) error

// Puller is used to sync data between to Databases
type Puller struct {
	waf WalkAddrs

	srcChunkStore nbs.NBSCompressedChunkStore
	sinkDBCS      chunks.ChunkStore
	hashes        hash.HashSet

	wr *PullTableFileWriter

	tempDir string
	pushLog *log.Logger

	statsCh chan Stats
	stats   *stats
}

// NewPuller creates a new Puller instance to do the syncing.  If a nil puller is returned without error that means
// that there is nothing to pull and the sinkDB is already up to date.
func NewPuller(
	ctx context.Context,
	tempDir string,
	targetFileSz uint64,
	srcCS, sinkCS chunks.ChunkStore,
	walkAddrs WalkAddrs,
	hashes []hash.Hash,
	statsCh chan Stats,
) (*Puller, error) {
	// Sanity Check
	hs := hash.NewHashSet(hashes...)
	missing, err := srcCS.HasMany(ctx, hs)
	if err != nil {
		return nil, err
	}
	if missing.Size() != 0 {
		return nil, errors.New("not found")
	}

	hs = hash.NewHashSet(hashes...)
	missing, err = sinkCS.HasMany(ctx, hs)
	if err != nil {
		return nil, err
	}
	if missing.Size() == 0 {
		return nil, ErrDBUpToDate
	}

	if srcCS.Version() != sinkCS.Version() {
		return nil, fmt.Errorf("cannot pull from src to sink; src version is %v and sink version is %v", srcCS.Version(), sinkCS.Version())
	}

	srcChunkStore, ok := srcCS.(nbs.NBSCompressedChunkStore)
	if !ok {
		return nil, ErrIncompatibleSourceChunkStore
	}

	// walkAddrs can be used for getAddrs on the AddTableFile call as well.
	getAddrs := func(c chunks.Chunk) chunks.InsertAddrsCb {
		return func(ctx context.Context, ins hash.HashSet, filter chunks.PendingRefExists) error {
			return walkAddrs(c, func(h hash.Hash, _ bool) error {
				if !filter(h) {
					ins.Insert(h)
				}
				return nil
			})
		}
	}

	wr := NewPullTableFileWriter(PullTableFileWriterConfig{
		ConcurrentUploads:    2,
		TargetFileSize:       targetFileSz,
		MaximumBufferedFiles: 8,
		TempDir:              tempDir,
		DestStore:            sinkCS.(chunks.TableFileStore),
		GetAddrs:             getAddrs,
	})

	p := &Puller{
		waf:           walkAddrs,
		srcChunkStore: srcChunkStore,
		sinkDBCS:      sinkCS,
		hashes:        hash.NewHashSet(hashes...),
		wr:            wr,
		tempDir:       tempDir,
		statsCh:       statsCh,
		stats: &stats{
			wrStatsGetter: wr.GetStats,
			fetcher:       &fetchStatsRecorder{},
		},
	}

	// Route chunk store debug logging into the push log. For a push, the sink
	// is the remote store; for a fetch, the source is. Wire both so either
	// operation's remote-side messages are captured.
	if lcs, ok := sinkCS.(chunks.LoggingChunkStore); ok {
		lcs.SetLogger(p)
	}
	if lcs, ok := srcCS.(chunks.LoggingChunkStore); ok {
		lcs.SetLogger(p)
	}

	return p, nil
}

func (p *Puller) Logf(fmt string, args ...interface{}) {
	if p.pushLog != nil {
		p.pushLog.Printf(fmt, args...)
	}
}

type readable interface {
	Reader() (io.ReadCloser, error)
	Remove() error
}

type tempTblFile struct {
	read        readable
	id          string
	contentHash []byte
	numChunks   int
	chunksLen   uint64
	contentLen  uint64
	pending     io.Closer
}

type countingReader struct {
	io.ReadCloser
	cnt *uint64
}

func (c countingReader) Read(p []byte) (int, error) {
	n, err := c.ReadCloser.Read(p)
	atomic.AddUint64(c.cnt, uint64(n))
	return n, err
}

func emitStats(s *stats, ch chan Stats, logf func(string, ...interface{})) (cancel func()) {
	done := make(chan struct{})
	var wg sync.WaitGroup
	wg.Add(2)
	cancel = func() {
		close(done)
		wg.Wait()
	}

	go func() {
		defer wg.Done()
		sampleduration := 100 * time.Millisecond
		samplesinsec := uint64((1 * time.Second) / sampleduration)
		weight := 0.1
		ticker := time.NewTicker(sampleduration)
		defer ticker.Stop()
		var lastSendBytes, lastFetchedBytes uint64
		for {
			select {
			case <-ticker.C:
				wrStats := s.wrStatsGetter()
				newSendBytes := wrStats.FinishedSendBytes
				newFetchedBytes := atomic.LoadUint64(&s.fetchedSourceBytes)
				sendBytesDiff := newSendBytes - lastSendBytes
				fetchedBytesDiff := newFetchedBytes - lastFetchedBytes

				newSendBPS := float64(sendBytesDiff * samplesinsec)
				newFetchedBPS := float64(fetchedBytesDiff * samplesinsec)

				curSendBPS := math.Float64frombits(atomic.LoadUint64(&s.sendBytesPerSec))
				curFetchedBPS := math.Float64frombits(atomic.LoadUint64(&s.fetchedSourceBytesPerSec))

				smoothedSendBPS := newSendBPS
				if curSendBPS != 0 {
					smoothedSendBPS = curSendBPS + weight*(newSendBPS-curSendBPS)
				}

				smoothedFetchBPS := newFetchedBPS
				if curFetchedBPS != 0 {
					smoothedFetchBPS = curFetchedBPS + weight*(newFetchedBPS-curFetchedBPS)
				}

				if smoothedSendBPS < 1 {
					smoothedSendBPS = 0
				}
				if smoothedFetchBPS < 1 {
					smoothedFetchBPS = 0
				}

				atomic.StoreUint64(&s.sendBytesPerSec, math.Float64bits(smoothedSendBPS))
				atomic.StoreUint64(&s.fetchedSourceBytesPerSec, math.Float64bits(smoothedFetchBPS))

				lastSendBytes = newSendBytes
				lastFetchedBytes = newFetchedBytes
			case <-done:
				return
			}
		}
	}()

	go func() {
		defer wg.Done()
		updateduration := 1 * time.Second
		ticker := time.NewTicker(updateduration)
		defer ticker.Stop()
		emit := func(st Stats) {
			if ch != nil {
				ch <- st
			}
			fs := s.fetcher.read()
			compBytes := atomic.LoadUint64(&s.fetchedSourceCompBytes)
			// Source side counters:
			//   chunks            - wanted chunks received / total requested
			//   uncompressed      - decompressed size of the wanted chunks
			//   compressed        - on-wire (snappy) size of just the wanted chunks
			//   wire              - bytes actually downloaded, including dark bytes
			//                       pulled in by range coalescing (wire - compressed
			//                       ~= coalescing overhead)
			//   fetch_rate        - uncompressed bytes/sec
			//   range_downloads   - completed coalesced range downloads (and retries)
			//   inflight/peak     - concurrent range downloads in flight
			// Sink side counters:
			//   sent/buffered     - table file bytes written to / queued for the dest
			//   send_rate         - sent bytes/sec
			logf("stats: source[chunks=%d/%d uncompressed=%dB compressed=%dB wire=%dB fetch_rate=%.0fB/s range_downloads=%d retries=%d inflight=%d peak=%d] sink[sent=%dB buffered=%dB send_rate=%.0fB/s]",
				st.FetchedSourceChunks, st.TotalSourceChunks,
				st.FetchedSourceBytes, compBytes, fs.DownloadedBytes, st.FetchedSourceBytesPerSec,
				fs.CompletedDownloads, fs.Retries, fs.InFlight, fs.PeakInFlight,
				st.FinishedSendBytes, st.BufferedSendBytes, st.SendBytesPerSec)
		}
		for {
			select {
			case <-ticker.C:
				emit(s.read())
			case <-done:
				emit(s.read())
				return
			}
		}
	}()

	return cancel
}

type stats struct {
	wrStatsGetter             func() PullTableFileWriterStats
	sendBytesPerSec           uint64
	totalSourceChunks         uint64
	fetchedSourceChunks       uint64
	fetchedSourceBytes        uint64
	fetchedSourceCompBytes    uint64
	fetchedSourceBytesPerSec  uint64
	sendBytesPerSecF          float64
	fetchedSourceBytesPerSecF float64

	// fetcher accumulates download stats reported by the source ChunkFetcher
	// (bytes fetched over the wire including dark bytes, download concurrency).
	fetcher *fetchStatsRecorder
}

type Stats struct {
	FinishedSendBytes uint64
	BufferedSendBytes uint64
	SendBytesPerSec   float64

	TotalSourceChunks        uint64
	FetchedSourceChunks      uint64
	FetchedSourceBytes       uint64
	FetchedSourceBytesPerSec float64
}

func (s *stats) read() Stats {
	wrStats := s.wrStatsGetter()

	var ret Stats
	ret.FinishedSendBytes = wrStats.FinishedSendBytes
	ret.BufferedSendBytes = wrStats.BufferedSendBytes
	ret.SendBytesPerSec = math.Float64frombits(atomic.LoadUint64(&s.sendBytesPerSec))
	ret.TotalSourceChunks = atomic.LoadUint64(&s.totalSourceChunks)
	ret.FetchedSourceChunks = atomic.LoadUint64(&s.fetchedSourceChunks)
	ret.FetchedSourceBytes = atomic.LoadUint64(&s.fetchedSourceBytes)
	ret.FetchedSourceBytesPerSec = math.Float64frombits(atomic.LoadUint64(&s.fetchedSourceBytesPerSec))
	return ret
}

func WithDiscardingStatsCh(cb func(statsCh chan Stats)) {
	statsCh := make(chan Stats)
	var wg sync.WaitGroup
	wg.Go(func() {
		for range statsCh {
		}
	})
	wg.Go(func() {
		defer close(statsCh)
		cb(statsCh)
	})
	wg.Wait()
}

// Pull executes the sync operation
func (p *Puller) Pull(ctx context.Context) error {
	// If push logging is enabled, create a unique log file for this pull
	// operation and close it when the pull completes.
	if dbg, ok := os.LookupEnv(dconfig.EnvPushLog); ok && strings.EqualFold(dbg, "true") {
		f, err := os.CreateTemp(p.tempDir, "push-*.log")
		if err == nil {
			p.pushLog = log.New(f, "", log.LstdFlags|log.Lmicroseconds)
			p.Logf("starting pull of %d hashes", p.hashes.Size())
			defer func() {
				p.Logf("pull complete")
				f.Close()
			}()
		}
	}

	// Run stats emission if anyone is consuming the stats channel or if we are
	// logging stats to the push log.
	if p.statsCh != nil || p.pushLog != nil {
		c := emitStats(p.stats, p.statsCh, p.Logf)
		defer c()
	}

	eg, ctx := errgroup.WithContext(ctx)

	rd := GetChunkFetcher(ctx, p.srcChunkStore, p.stats.fetcher)

	const batchSize = 64 * 1024
	tracker := NewPullChunkTracker(TrackerConfig{
		BatchSize: batchSize,
		HasManyer: p.sinkDBCS,
	})

	eg.Go(func() error {
		return tracker.Run(ctx, p.hashes)
	})

	eg.Go(func() error {
		return p.wr.Run(ctx)
	})

	// One thread calls ChunkFetcher.Get on each batch.
	eg.Go(func() error {
		for {
			toFetch, hasMore, err := tracker.GetChunksToFetch(ctx)
			if err != nil {
				return err
			}
			if !hasMore {
				return rd.CloseSend()
			}

			atomic.AddUint64(&p.stats.totalSourceChunks, uint64(len(toFetch)))
			err = rd.Get(ctx, toFetch)
			if err != nil {
				return err
			}
		}
	})

	// One thread reads the received chunks, walks their addresses and writes them to the table file.
	eg.Go(func() (err error) {
		// Closing the tracker belongs here, because this is
		// the goroutine that writes to the tracker state with
		// |Seen| and |TickProcessed| calls.  Once this thread
		// finishes, there will be no further writes to the
		// tracker state.
		defer tracker.Close()
		defer func() {
			cerr := rd.Close()
			err = errors.Join(err, cerr)
		}()
		for {
			cChk, err := rd.Recv(ctx)
			if err == io.EOF {
				// This means the requesting thread
				// successfully saw all chunk addresses and
				// called CloseSend and that all requested
				// chunks were successfully delivered to this
				// thread. Calling wr.Close() here will block
				// on uploading any table files and will write
				// the new table files to the destination's
				// manifest.
				p.wr.Close()
				return nil
			}
			if err != nil {
				return err
			}
			if cChk.IsGhost() {
				return fmt.Errorf("attempted to push or pull ghost chunk: %w", nbs.ErrGhostChunkRequested)
			}
			if cChk.IsEmpty() {
				return errors.New("failed to get all chunks.")
			}

			atomic.AddUint64(&p.stats.fetchedSourceChunks, uint64(1))
			// CompressedSize is the on-the-wire (snappy-compressed) size of
			// just this wanted chunk, excluding any dark bytes pulled in by
			// range coalescing (those are tracked by the fetcher recorder).
			atomic.AddUint64(&p.stats.fetchedSourceCompBytes, uint64(cChk.CompressedSize()))

			chnk, err := cChk.ToChunk()
			if err != nil {
				return err
			}

			atomic.AddUint64(&p.stats.fetchedSourceBytes, uint64(len(chnk.Data())))

			err = p.waf(chnk, func(h hash.Hash, _ bool) error {
				tracker.Seen(ctx, h)
				return nil
			})
			if err != nil {
				return err
			}
			tracker.TickProcessed(ctx)

			err = p.wr.AddToChunker(ctx, cChk)
			if err != nil {
				return err
			}
		}
	})

	return eg.Wait()
}