// Copyright (C) MongoDB, Inc. 2014-present.
//
// 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

package archive

import (
	"bytes"
	"fmt"
	"hash"
	"hash/crc64"
	"io"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/mongodb/mongo-tools/common"
	"github.com/mongodb/mongo-tools/common/db"
	"github.com/mongodb/mongo-tools/common/intents"
	"github.com/mongodb/mongo-tools/common/log"
	"github.com/mongodb/mongo-tools/common/util"
	"go.mongodb.org/mongo-driver/v2/bson"
)

// DemuxOut is a Demultiplexer output consumer
// The Write() and End() occur in the same thread as the Demultiplexer runs in.
type DemuxOut interface {
	Write([]byte) (int, error)
	End()
	Sum64() (uint64, bool)
}

const (
	NamespaceUnopened = iota
	NamespaceOpened
	NamespaceClosed
)

// Demultiplexer implements Parser.
type Demultiplexer struct {
	In io.Reader
	//TODO wrap up these three into a structure
	outs             map[string]DemuxOut
	lengths          map[string]int64
	currentNamespace string

	// NamespaceChan is used to send a namespace to a consumer of namespaces.
	NamespaceChan chan string

	// NamespaceErrorChan is used to receive an error or nil from a namespace
	// consumer immediately after each namespace is sent and validated.
	NamespaceErrorChan chan error

	NamespaceStatus map[string]int
	IsAtlasProxy    bool
}

// CreateDemux builds a Demultiplexer for an archive. sourceVersion must be the
// server version that *produced* the archive (from the archive prelude), not the
// destination server version. The physical namespace that timeseries data is
// stored under in the archive is determined by the source server: pre-8.3 sources
// store it under "system.buckets.<coll>", while viewless-timeseries sources (8.3+)
// store it under "<coll>". This must match how PreludeExplorer derives the
// receiver namespace, otherwise the demux listens on the wrong namespace.
func CreateDemux(
	sourceVersion db.Version,
	namespaceMetadatas []*CollectionMetadata,
	in io.Reader,
	isAtlasProxy bool,
) *Demultiplexer {
	demux := &Demultiplexer{
		NamespaceStatus: make(map[string]int),
		In:              in,
		IsAtlasProxy:    isAtlasProxy,
	}
	for _, cm := range namespaceMetadatas {
		// For atlas proxy archive restores, ignore collections from the admin DB.
		if isAtlasProxy && cm.Database == "admin" {
			continue
		}

		var ns string
		if cm.Type == "timeseries" && !sourceVersion.SupportsRawData() {
			// Sources older than 8.3 store timeseries data under system.buckets;
			// 8.3+ sources use viewless timeseries stored under the collection name.
			ns = cm.Database + "." + common.TimeseriesBucketPrefix + cm.Collection
		} else {
			ns = cm.Database + "." + cm.Collection
		}
		demux.NamespaceStatus[ns] = NamespaceUnopened
	}
	return demux
}

// Run creates and runs a parser with the Demultiplexer as a consumer.
func (demux *Demultiplexer) Run() error {
	parser := Parser{In: demux.In}
	err := parser.ReadAllBlocks(demux)
	if len(demux.outs) > 0 {
		log.Logvf(log.Always, "demux finishing when there are still outs (%v)", len(demux.outs))
	}

	log.Logvf(log.DebugLow, "demux finishing (err:%v)", err)
	return err
}

type demuxError struct {
	Err error
	Msg string
}

// Error is part of the Error interface. It formats a demuxError for human readability.
func (pe *demuxError) Error() string {
	err := fmt.Sprintf("error demultiplexing archive; %v", pe.Msg)
	if pe.Err != nil {
		err = fmt.Sprintf("%v ( %v )", err, pe.Err)
	}
	return err
}

// newError creates a demuxError with just a message.
func newError(msg string) error {
	return &demuxError{
		Msg: msg,
	}
}

// newWrappedError creates a demuxError with a message as well as an underlying cause error.
func newWrappedError(msg string, err error) error {
	return &demuxError{
		Err: err,
		Msg: msg,
	}
}

// HeaderBSON is part of the ParserConsumer interface and receives headers from parser.
// Its main role is to implement opens and EOFs of the embedded stream.
func (demux *Demultiplexer) HeaderBSON(buf []byte) error {
	colHeader := NamespaceHeader{}
	err := bson.Unmarshal(buf, &colHeader)
	if err != nil {
		return newWrappedError("header bson doesn't unmarshal as a collection header", err)
	}
	log.Logvf(log.DebugHigh, "demux namespaceHeader: %v", colHeader)
	if colHeader.Collection == "" {
		return newError("collection header is missing a Collection")
	}
	demux.currentNamespace = colHeader.Database + "." + colHeader.Collection

	// For atlas proxy archive restores, ignore collections from the admin DB.
	if demux.IsAtlasProxy && colHeader.Database == "admin" {
		return nil
	}

	if _, ok := demux.outs[demux.currentNamespace]; !ok {
		if demux.NamespaceStatus[demux.currentNamespace] != NamespaceUnopened {
			return newError("namespace header for already opened namespace")
		}
		demux.NamespaceStatus[demux.currentNamespace] = NamespaceOpened
		if demux.NamespaceChan != nil {
			demux.NamespaceChan <- demux.currentNamespace
			err := <-demux.NamespaceErrorChan
			if err == io.EOF {
				// if the Prioritizer sends us back an io.EOF then it's telling us that
				// it's finishing and doesn't need any more namespace announcements.
				close(demux.NamespaceChan)
				demux.NamespaceChan = nil
				return nil
			}
			if err != nil {
				return newWrappedError("failed arranging a consumer for new namespace", err)
			}
		}
	}
	if colHeader.EOF {
		if rcr, ok := demux.outs[demux.currentNamespace].(*RegularCollectionReceiver); ok {
			rcr.err = io.EOF
		}
		demux.outs[demux.currentNamespace].End()
		demux.NamespaceStatus[demux.currentNamespace] = NamespaceClosed
		length := demux.lengths[demux.currentNamespace]
		crc, ok := demux.outs[demux.currentNamespace].Sum64()
		if ok {
			if crc != colHeader.CRC {
				return fmt.Errorf("CRC mismatch for namespace %#q, %v!=%v",
					demux.currentNamespace,
					crc,
					colHeader.CRC,
				)
			}
			log.Logvf(log.DebugHigh,
				"demux checksum for namespace %#q is correct (%v), %v bytes",
				demux.currentNamespace, crc, length)
		} else {
			log.Logvf(log.DebugHigh,
				"demux checksum for namespace %#q was not calculated.",
				demux.currentNamespace)
		}
		delete(demux.outs, demux.currentNamespace)
		delete(demux.lengths, demux.currentNamespace)
		// in case we get a BSONBody with this block,
		// we want to ensure that this causes an error
		demux.currentNamespace = ""
	}
	return nil
}

// End is part of the ParserConsumer interface and receives the end of archive notification.
func (demux *Demultiplexer) End() error {
	log.Logvf(log.DebugHigh, "demux End")
	var err error
	if len(demux.outs) != 0 {
		openNss := []string{}
		for ns := range demux.outs {
			openNss = append(openNss, ns)
			if rcr, ok := demux.outs[ns].(*RegularCollectionReceiver); ok {
				rcr.err = newError("archive io error")
			}
			demux.outs[ns].End()
		}
		err = newError(
			fmt.Sprintf(
				"archive finished but contained files were unfinished (%v)",
				util.QuoteAndJoin(openNss, ", "),
			),
		)
	} else {
		for ns, status := range demux.NamespaceStatus {
			if status != NamespaceClosed {
				err = newError(fmt.Sprintf("archive finished before all collections were seen (%#q)", ns))
			}
		}
	}

	if demux.NamespaceChan != nil {
		close(demux.NamespaceChan)
	}

	return err
}

// BodyBSON is part of the ParserConsumer interface and receives BSON bodies from the parser.
// Its main role is to dispatch the body to the Read() function of the current DemuxOut.
func (demux *Demultiplexer) BodyBSON(buf []byte) error {
	if demux.currentNamespace == "" {
		return newError("collection data without a collection header")
	}

	// For atlas proxy archive restores, ignore collections from the admin DB.
	if demux.IsAtlasProxy && strings.HasPrefix(demux.currentNamespace, "admin.") {
		return nil
	}

	demux.lengths[demux.currentNamespace] += int64(len(buf))

	out, ok := demux.outs[demux.currentNamespace]
	if !ok {
		return newError("no demux consumer currently consuming namespace " + demux.currentNamespace)
	}
	_, err := out.Write(buf)
	return err
}

// Open installs the DemuxOut as the handler for data for the namespace ns.
func (demux *Demultiplexer) Open(ns string, out DemuxOut) {
	// In the current implementation where this is either called before the demultiplexing is running
	// or while the demutiplexer is inside of the NamespaceChan NamespaceErrorChan conversation
	// I think that we don't need to lock outs, but I suspect that if the implementation changes
	// we may need to lock when outs is accessed
	log.Logvf(log.DebugHigh, "demux Open for %#q", ns)
	if demux.outs == nil {
		demux.outs = make(map[string]DemuxOut)
		demux.lengths = make(map[string]int64)
	}
	demux.outs[ns] = out
	demux.lengths[ns] = 0
}

// RegularCollectionReceiver implements the intents.file interface.
type RegularCollectionReceiver struct {
	pos              int64 // updated atomically, aligned at the beginning of the struct
	readLenChan      chan int
	readBufChan      chan []byte
	Intent           *intents.Intent
	Origin           string
	Demux            *Demultiplexer
	partialReadArray []byte
	partialReadBuf   []byte
	hash             hash.Hash64
	closeOnce        sync.Once
	endOnce          sync.Once
	openOnce         sync.Once
	err              error
}

func (receiver *RegularCollectionReceiver) Sum64() (uint64, bool) {
	return receiver.hash.Sum64(), true
}

// Read() runs in the restoring goroutine.
func (receiver *RegularCollectionReceiver) Read(r []byte) (int, error) {
	if receiver.partialReadBuf != nil && len(receiver.partialReadBuf) > 0 {
		wLen := len(receiver.partialReadBuf)
		copyLen := copy(r, receiver.partialReadBuf)
		if wLen == copyLen {
			receiver.partialReadBuf = nil
		} else {
			receiver.partialReadBuf = receiver.partialReadBuf[copyLen:]
		}
		atomic.AddInt64(&receiver.pos, int64(copyLen))
		return copyLen, nil
	}
	// Since we're the "reader" here, not the "writer" we need to start with a read, in case the chan is closed
	wLen, ok := <-receiver.readLenChan
	if !ok {
		return 0, receiver.err
	}
	if wLen > db.MaxBSONSize {
		return 0, fmt.Errorf("incoming buffer size is too big %v", wLen)
	}
	rLen := len(r)
	if wLen > rLen {
		// if the incoming write size is larger then the incoming read buffer then we need to accept
		// the write in a larger buffer, fill the read buffer, then cache the remainder
		receiver.partialReadBuf = receiver.partialReadArray[:wLen]
		receiver.readBufChan <- receiver.partialReadBuf
		writtenLength := <-receiver.readLenChan
		if wLen != writtenLength {
			return 0, fmt.Errorf("regularCollectionReceiver didn't send what it said it would")
		}
		receiver.hash.Write(receiver.partialReadBuf)
		copy(r, receiver.partialReadBuf)
		receiver.partialReadBuf = receiver.partialReadBuf[rLen:]
		atomic.AddInt64(&receiver.pos, int64(rLen))
		return rLen, nil
	}
	// Send the read buff to the BodyBSON ParserConsumer to fill
	receiver.readBufChan <- r
	// Receiver the wLen of data written
	wLen = <-receiver.readLenChan
	receiver.hash.Write(r[:wLen])
	atomic.AddInt64(&receiver.pos, int64(wLen))
	return wLen, nil
}

func (receiver *RegularCollectionReceiver) Pos() int64 {
	return atomic.LoadInt64(&receiver.pos)
}

// Open is part of the intents.file interface.  It creates the chan's in the
// RegularCollectionReceiver and adds the RegularCollectionReceiver to the set of
// RegularCollectionReceivers in the demultiplexer.
func (receiver *RegularCollectionReceiver) Open() error {
	// TODO move this implementation to some non intents.file method, to be called from prioritizer.Get
	// So that we don't have to enable this double open stuff.
	// Currently the open needs to finish before the prioritizer.Get finishes, so we open the intents.file
	// in prioritizer.Get even though it's going to get opened again in DumpIntent.
	receiver.openOnce.Do(func() {
		receiver.readLenChan = make(chan int)
		receiver.readBufChan = make(chan []byte)
		receiver.hash = crc64.New(crc64.MakeTable(crc64.ECMA))
		receiver.Demux.Open(receiver.Origin, receiver)
	})
	return nil
}

func (receiver *RegularCollectionReceiver) TakeIOBuffer(ioBuf []byte) {
	receiver.partialReadArray = ioBuf

}
func (receiver *RegularCollectionReceiver) ReleaseIOBuffer() {
	receiver.partialReadArray = nil
}

// Write is part of the DemuxOut interface.
func (receiver *RegularCollectionReceiver) Write(buf []byte) (int, error) {
	//  As a writer, we need to write first, so that the reader can properly detect EOF
	//  Additionally, the reader needs to know the write size, so that it can give us a
	//  properly sized buffer. Sending the incoming buffersize fills both of these needs.
	//  However, it's possible the intent has been closed already so the reader
	//  won't be reading the channel.  To avoid deadlock, we check for closed
	//  readBufChan.
	select {
	case receiver.readLenChan <- len(buf):
		// do nothing other than send
	case <-receiver.readBufChan:
		// Will only receive this if closed
		return 0, errInterrupted
	}

	// Receive from the reader a buffer to put the bytes into
	readBuf := <-receiver.readBufChan
	if len(readBuf) < len(buf) {
		return 0, fmt.Errorf("readbuf is not large enough for incoming BodyBSON (%v<%v)",
			len(readBuf), len(buf))
	}
	copy(readBuf, buf)
	// Send back the length of the data copied in to the buffer
	receiver.readLenChan <- len(buf)
	return len(buf), nil
}

// Close is part of the DemuxOut as well as the intents.file interface. It only closes the readLenChan, as that is what will
// cause the RegularCollectionReceiver.Read() to receive EOF
// Close will get called twice, once in the demultiplexer, and again when the restore goroutine is done with its intent.file.
func (receiver *RegularCollectionReceiver) Close() error {
	// Close must be idempotent and repeat channel closes panic; only do once.
	receiver.closeOnce.Do(func() {
		close(receiver.readBufChan)
	})
	return nil
}

// End signals to any waiting readers that there is nothing more to read, then
// it waits on the readBufChan, which is closed by the reader-side Close()
// method.
func (receiver *RegularCollectionReceiver) End() {
	// To keep this idempotent, close the channel only once.
	receiver.endOnce.Do(func() {
		close(receiver.readLenChan)
	})
	<-receiver.readBufChan
}

// SpecialCollectionCache implements both DemuxOut as well as intents.file.
type SpecialCollectionCache struct {
	pos    int64 // updated atomically, aligned at the beginning of the struct
	Intent *intents.Intent
	Demux  *Demultiplexer
	buf    bytes.Buffer
	hash   hash.Hash64
}

func NewSpecialCollectionCache(
	intent *intents.Intent,
	demux *Demultiplexer,
) *SpecialCollectionCache {
	return &SpecialCollectionCache{
		Intent: intent,
		Demux:  demux,
		hash:   crc64.New(crc64.MakeTable(crc64.ECMA)),
	}
}

// Open is part of the both interfaces, and it does nothing.
func (cache *SpecialCollectionCache) Open() error {
	return nil
}

// Close is part of the intents.file interface, and does nothing.
func (cache *SpecialCollectionCache) Close() error {
	return nil
}

// End indicates we've read all there is to read, so
// we update the intent size based on the final length.
func (cache *SpecialCollectionCache) End() {
	cache.Intent.Size = int64(cache.buf.Len())
}

func (cache *SpecialCollectionCache) Read(p []byte) (int, error) {
	n, err := cache.buf.Read(p)
	atomic.AddInt64(&cache.pos, int64(n))
	return n, err
}

func (cache *SpecialCollectionCache) Pos() int64 {
	return atomic.LoadInt64(&cache.pos)
}

func (cache *SpecialCollectionCache) Write(b []byte) (int, error) {
	// Writes to the hash never return an error.
	cache.hash.Write(b)
	return cache.buf.Write(b)
}

func (cache *SpecialCollectionCache) Sum64() (uint64, bool) {
	return cache.hash.Sum64(), true
}

// MutedCollection implements both DemuxOut as well as intents.file. It serves as a way to
// let the demutiplexer ignore certain embedded streams.
type MutedCollection struct {
	Intent *intents.Intent
	Demux  *Demultiplexer
}

// Read is part of the intents.file interface, and does nothing.
func (*MutedCollection) Read([]byte) (int, error) {
	// Read is part of the intents.file interface, and does nothing
	return 0, io.EOF
}

// Write is part of the intents.file interface, and does nothing.
func (*MutedCollection) Write(b []byte) (int, error) {
	return len(b), nil
}

// Close is part of the intents.file interface, and does nothing.
func (*MutedCollection) Close() error {
	return nil
}

// End is part of the DemuxOut interface and does nothing.
func (*MutedCollection) End() {}

// Open is part of the intents.file interface, and does nothing.
func (*MutedCollection) Open() error {
	return nil
}

// Sum64 is part of the DemuxOut interface.
func (*MutedCollection) Sum64() (uint64, bool) {
	return 0, false
}

//===== Archive Manager Prioritizer =====

// NewPrioritizer creates a new Prioritizer and hooks up its Namespace channels to the ones in demux.
func (demux *Demultiplexer) NewPrioritizer(mgr *intents.Manager) *Prioritizer {
	return &Prioritizer{
		NamespaceChan:      demux.NamespaceChan,
		NamespaceErrorChan: demux.NamespaceErrorChan,
		mgr:                mgr,
	}
}

// Prioritizer is a completely reactive prioritizer
// Intents are handed out as they arrive in the archive.
type Prioritizer struct {
	NamespaceChan      <-chan string
	NamespaceErrorChan chan<- error
	mgr                *intents.Manager
}

// Get waits for a new namespace from the NamespaceChan, and returns a Intent found for it.
func (prioritizer *Prioritizer) Get() *intents.Intent {
	namespace, ok := <-prioritizer.NamespaceChan
	if !ok {
		return nil
	}
	destDB, destC := util.SplitNamespace(namespace)
	namespace = destDB + "." + strings.TrimPrefix(destC, common.TimeseriesBucketPrefix)
	intent := prioritizer.mgr.IntentForNamespace(namespace)
	if intent == nil {
		prioritizer.NamespaceErrorChan <- fmt.Errorf("no intent for namespace %#q", namespace)
	} else {
		if intent.BSONFile != nil {
			if err := intent.BSONFile.Open(); err != nil {
				prioritizer.NamespaceErrorChan <- fmt.Errorf("error opening BSON file for %#q: %v", intent.Namespace(), err)
			}
		}
		if intent.IsOplog() {
			// once we see the oplog we
			// cause the RestoreIntents to finish because we don't
			// want RestoreIntents to restore the oplog
			prioritizer.NamespaceErrorChan <- io.EOF
			return nil
		}
		prioritizer.NamespaceErrorChan <- nil
	}
	return intent
}

// Finish is part of the IntentPrioritizer interface, and does nothing.
func (prioritizer *Prioritizer) Finish(*intents.Intent) {
	// no-op
	return
}
