// 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"
	"hash"
	"hash/crc32"
	"io"
	"os"
	"sync"
	"testing"

	"github.com/mongodb/mongo-tools/common/db"
	"github.com/mongodb/mongo-tools/common/intents"
	"github.com/mongodb/mongo-tools/common/testtype"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.mongodb.org/mongo-driver/v2/bson"
	"golang.org/x/sync/errgroup"
)

const testDocCount = 10000

var testIntents = []*intents.Intent{
	{
		ServerVersion: db.Version{7, 0, 0},
		DB:            "foo",
		C:             "bar",
		Location:      "foo.bar",
	},
	{
		ServerVersion: db.Version{7, 0, 0},
		DB:            "ding",
		C:             "bats",
		Location:      "ding.bats",
	},
	{
		ServerVersion: db.Version{7, 0, 0},
		DB:            "flim",
		C:             "flam.fooey",
		Location:      "flim.flam.fooey",
	},
	{
		ServerVersion: db.Version{7, 0, 0},
		DB:            "crow",
		C:             "bar",
		Location:      "crow.bar",
	},
}

type testDoc struct {
	Bar int
	Baz string
}

type closingBuffer struct {
	bytes.Buffer
}

func (*closingBuffer) Close() error {
	return nil
}

type testNotifier struct{}

func (n *testNotifier) Notify() {}

func TestBasicMux(t *testing.T) {
	testtype.SkipUnlessTestType(t, testtype.UnitTestType)

	var err error

	buf := &closingBuffer{bytes.Buffer{}}

	mux := NewMultiplexer(buf, new(testNotifier))
	muxIns := map[string]*MuxIn{}

	inChecksum := map[string]hash.Hash{}
	inLengths := map[string]*int{}
	outChecksum := map[string]hash.Hash{}
	outLengths := map[string]*int{}

	// To confirm that what we multiplex is the same as what we demultiplex, we
	// create input and output hashes for each namespace. After we finish
	// multiplexing and demultiplexing we will compare all of the CRCs for each
	// namespace
	errChan := make(chan error)
	makeIns(testIntents, mux, inChecksum, muxIns, inLengths, errChan)

	t.Run("document multiplexing", func(t *testing.T) {
		go mux.Run()

		for range testIntents {
			err := <-errChan
			require.NoError(t, err)
		}
		close(mux.Control)
		err = <-mux.Completed
		require.NoError(t, err)
	})

	t.Run("document demultiplexing", func(t *testing.T) {
		demux := &Demultiplexer{
			In:              buf,
			NamespaceStatus: make(map[string]int),
		}
		demuxOuts := map[string]*RegularCollectionReceiver{}

		errChan := make(chan error)
		makeOuts(t, testIntents, demux, outChecksum, demuxOuts, outLengths, errChan)

		err := demux.Run()
		require.NoError(t, err)

		for range testIntents {
			err := <-errChan
			require.NoError(t, err)
		}
		for _, dbc := range testIntents {
			ns := dbc.Namespace()
			assert.Equal(t, *inLengths[ns], *outLengths[ns])
			inSum := inChecksum[ns].Sum([]byte{})
			outSum := outChecksum[ns].Sum([]byte{})
			assert.Equal(t, inSum, outSum)
		}
	})

}

func TestParallelMuxOverPipe(t *testing.T) {
	testtype.SkipUnlessTestType(t, testtype.UnitTestType)

	readPipe, writePipe, err := os.Pipe()
	require.NoError(t, err)

	mux := NewMultiplexer(writePipe, new(testNotifier))
	muxIns := map[string]*MuxIn{}

	demux := &Demultiplexer{
		In:              readPipe,
		NamespaceStatus: make(map[string]int),
	}
	demuxOuts := map[string]*RegularCollectionReceiver{}

	inChecksum := map[string]hash.Hash{}
	inLengths := map[string]*int{}

	outChecksum := map[string]hash.Hash{}
	outLengths := map[string]*int{}

	writeErrChan := make(chan error)
	readErrChan := make(chan error)

	makeIns(testIntents, mux, inChecksum, muxIns, inLengths, writeErrChan)
	makeOuts(t, testIntents, demux, outChecksum, demuxOuts, outLengths, readErrChan)

	var eg errgroup.Group

	eg.Go(demux.Run)
	eg.Go(func() error {
		mux.Run()
		return nil
	})

	for range testIntents {
		err := <-writeErrChan
		require.NoError(t, err)
		err = <-readErrChan
		require.NoError(t, err)
	}
	close(mux.Control)
	muxErr := <-mux.Completed
	require.NoError(t, muxErr)

	for _, dbc := range testIntents {
		ns := dbc.Namespace()
		assert.Equal(t, *inLengths[ns], *outLengths[ns])
		inSum := inChecksum[ns].Sum([]byte{})
		outSum := outChecksum[ns].Sum([]byte{})
		assert.Equal(t, inSum, outSum)
	}

	demuxErr := eg.Wait()
	require.NoError(t, demuxErr)
}

func makeIns(
	testIntents []*intents.Intent,
	mux *Multiplexer,
	inChecksum map[string]hash.Hash,
	muxIns map[string]*MuxIn,
	inLengths map[string]*int,
	errCh chan<- error,
) {
	for index, dbc := range testIntents {
		ns := dbc.Namespace()
		sum := crc32.NewIEEE()
		muxIn := &MuxIn{Intent: dbc, Mux: mux}
		inLength := 0

		inChecksum[ns] = sum
		muxIns[ns] = muxIn
		inLengths[ns] = &inLength

		go func(index int) {
			err := muxIn.Open()
			if err != nil {
				errCh <- err
				return
			}
			staticBSONBuf := make([]byte, db.MaxBSONSize)
			for i := 0; i < testDocCount; i++ {
				bsonBytes, _ := bson.Marshal(testDoc{Bar: index * i, Baz: ns})
				bsonBuf := staticBSONBuf[:len(bsonBytes)]
				copy(bsonBuf, bsonBytes)
				_, err := muxIn.Write(bsonBuf)
				if err != nil {
					errCh <- err
					return
				}
				sum.Write(bsonBuf)
				inLength += len(bsonBuf)
			}
			err = muxIn.Close()
			errCh <- err
		}(index)
	}
}

func makeOuts(
	t *testing.T,
	testIntents []*intents.Intent,
	demux *Demultiplexer,
	outChecksum map[string]hash.Hash,
	demuxOuts map[string]*RegularCollectionReceiver,
	outLengths map[string]*int,
	errCh chan<- error,
) {
	for _, dbc := range testIntents {
		ns := dbc.Namespace()
		sum := crc32.NewIEEE()
		muxOut := &RegularCollectionReceiver{
			Intent: dbc,
			Demux:  demux,
			Origin: ns,
		}
		outLength := 0

		outChecksum[ns] = sum
		demuxOuts[ns] = muxOut
		outLengths[ns] = &outLength

		err := demuxOuts[ns].Open()
		require.NoError(t, err)

		go func() {
			bs := make([]byte, db.MaxBSONSize)
			var err error
			for {
				var length int
				length, err = muxOut.Read(bs)
				if err != nil {
					muxOut.Close()
					break
				}
				sum.Write(bs[:length])
				outLength += len(bs[:length])
			}
			if err == io.EOF {
				err = nil
			}
			errCh <- err
		}()
	}
}

func buildSingleIntentArchive(t *testing.T, singleIntent *intents.Intent) *closingBuffer {
	var err error

	buf := &closingBuffer{bytes.Buffer{}}

	mux := NewMultiplexer(buf, new(testNotifier))
	muxIns := map[string]*MuxIn{}

	inChecksum := map[string]hash.Hash{}
	inLengths := map[string]*int{}
	errChan := make(chan error)

	makeIns([]*intents.Intent{singleIntent}, mux, inChecksum, muxIns, inLengths, errChan)

	go mux.Run()

	err = <-errChan
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	close(mux.Control)
	err = <-mux.Completed
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	return buf
}

func TestTOOLS1826(t *testing.T) {
	testtype.SkipUnlessTestType(t, testtype.UnitTestType)
	require := require.New(t)

	singleIntent := testIntents[0]

	demux := &Demultiplexer{
		In:              buildSingleIntentArchive(t, singleIntent),
		NamespaceStatus: make(map[string]int),
	}

	muxOut := &RegularCollectionReceiver{
		Intent: singleIntent,
		Demux:  demux,
		Origin: singleIntent.Namespace(),
	}
	require.NoError(muxOut.Open())

	var demuxErr error
	wg := &sync.WaitGroup{}
	wg.Add(1)
	go func() {
		defer wg.Done()
		demuxErr = demux.Run()
	}()

	// Closing the receiver before reading shouldn't panic.
	muxOut.Close()

	wg.Wait()
	require.ErrorIs(demuxErr, errInterrupted)
}

func TestTOOLS2403(t *testing.T) {
	testtype.SkipUnlessTestType(t, testtype.UnitTestType)
	require := require.New(t)

	singleIntent := testIntents[0]

	demux := &Demultiplexer{
		In:              buildSingleIntentArchive(t, singleIntent),
		NamespaceStatus: make(map[string]int),
	}

	muxOut := &RegularCollectionReceiver{
		Intent: singleIntent,
		Demux:  demux,
		Origin: singleIntent.Namespace(),
	}
	require.NoError(muxOut.Open())

	var demuxErr error
	wg := &sync.WaitGroup{}
	wg.Add(1)
	go func() {
		defer wg.Done()
		demuxErr = demux.Run()
	}()

	// Read all the documents, but don't read past into EOF.
	bs := make([]byte, db.MaxBSONSize)
	for i := 0; i < testDocCount; i++ {
		_, err := muxOut.Read(bs)
		if err != nil {
			t.Fatal(err)
		}
	}

	// Closing the intent before reading EOF should not deadlock.
	muxOut.Close()

	wg.Wait()
	require.NoError(demuxErr)

	return
}
