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

import (
	"testing"

	"github.com/mongodb/mongo-tools/common/log"
	"github.com/mongodb/mongo-tools/common/options"
	"github.com/mongodb/mongo-tools/common/testtype"
	. "github.com/smartystreets/goconvey/convey"
	"go.mongodb.org/mongo-driver/v2/bson"
	"golang.org/x/sync/errgroup"
)

func init() {
	log.SetVerbosity(&options.Verbosity{
		VLevel: 4,
	})
}

var (
	index         = uint64(0)
	csvConverters = []CSVConverter{
		{
			colSpecs: []ColumnSpec{
				{"field1", new(FieldAutoParser), pgAutoCast, "auto", []string{"field1"}},
				{"field2", new(FieldAutoParser), pgAutoCast, "auto", []string{"field2"}},
				{"field3", new(FieldAutoParser), pgAutoCast, "auto", []string{"field3"}},
			},
			data:  []string{"a", "b", "c"},
			index: index,
		},
		{
			colSpecs: []ColumnSpec{
				{"field4", new(FieldAutoParser), pgAutoCast, "auto", []string{"field4"}},
				{"field5", new(FieldAutoParser), pgAutoCast, "auto", []string{"field5"}},
				{"field6", new(FieldAutoParser), pgAutoCast, "auto", []string{"field6"}},
			},
			data:  []string{"d", "e", "f"},
			index: index,
		},
		{
			colSpecs: []ColumnSpec{
				{"field7", new(FieldAutoParser), pgAutoCast, "auto", []string{"field7"}},
				{"field8", new(FieldAutoParser), pgAutoCast, "auto", []string{"field8"}},
				{"field9", new(FieldAutoParser), pgAutoCast, "auto", []string{"field9"}},
			},
			data:  []string{"d", "e", "f"},
			index: index,
		},
		{
			colSpecs: []ColumnSpec{
				{"field10", new(FieldAutoParser), pgAutoCast, "auto", []string{"field10"}},
				{"field11", new(FieldAutoParser), pgAutoCast, "auto", []string{"field11"}},
				{"field12", new(FieldAutoParser), pgAutoCast, "auto", []string{"field12"}},
			},
			data:  []string{"d", "e", "f"},
			index: index,
		},
		{
			colSpecs: []ColumnSpec{
				{"field13", new(FieldAutoParser), pgAutoCast, "auto", []string{"field13"}},
				{"field14", new(FieldAutoParser), pgAutoCast, "auto", []string{"field14"}},
				{"field15", new(FieldAutoParser), pgAutoCast, "auto", []string{"field15"}},
			},
			data:  []string{"d", "e", "f"},
			index: index,
		},
	}
	expectedDocuments = []bson.D{
		{
			{"field1", "a"},
			{"field2", "b"},
			{"field3", "c"},
		}, {
			{"field4", "d"},
			{"field5", "e"},
			{"field6", "f"},
		}, {
			{"field7", "d"},
			{"field8", "e"},
			{"field9", "f"},
		}, {
			{"field10", "d"},
			{"field11", "e"},
			{"field12", "f"},
		}, {
			{"field13", "d"},
			{"field14", "e"},
			{"field15", "f"},
		},
	}
)

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

	Convey("Given an import input, in validating the headers", t, func() {
		Convey("if the fields contain '..', an error should be thrown", func() {
			So(validateFields([]string{"a..a"}, false), ShouldNotBeNil)
		})
		Convey("if the fields start/end in a '.', an error should be thrown", func() {
			So(validateFields([]string{".a"}, false), ShouldNotBeNil)
			So(validateFields([]string{"a."}, false), ShouldNotBeNil)
		})
		Convey("if the fields start in a '$', an error should be thrown", func() {
			So(validateFields([]string{"$.a"}, false), ShouldNotBeNil)
			So(validateFields([]string{"$"}, false), ShouldNotBeNil)
			So(validateFields([]string{"$a"}, false), ShouldNotBeNil)
			So(validateFields([]string{"a$a"}, false), ShouldBeNil)
		})
		Convey("if the fields collide, an error should be thrown", func() {
			So(validateFields([]string{"a", "a.a"}, false), ShouldNotBeNil)
			So(validateFields([]string{"a", "a.ba", "b.a"}, false), ShouldNotBeNil)
			So(validateFields([]string{"a", "a.ba", "b.a"}, false), ShouldNotBeNil)
			So(validateFields([]string{"a", "a.b.c"}, false), ShouldNotBeNil)
		})
		Convey("if the fields don't collide, no error should be thrown", func() {
			So(validateFields([]string{"a", "aa"}, false), ShouldBeNil)
			So(validateFields([]string{"a", "aa", "b.a", "b.c"}, false), ShouldBeNil)
			So(validateFields([]string{"a", "ba", "ab", "b.a"}, false), ShouldBeNil)
			So(validateFields([]string{"a", "ba", "ab", "b.a", "b.c.d"}, false), ShouldBeNil)
			So(validateFields([]string{"a", "ab.c"}, false), ShouldBeNil)
		})
		Convey("if the fields contain the same keys, an error should be thrown", func() {
			So(validateFields([]string{"a", "ba", "a"}, false), ShouldNotBeNil)
		})
	})
}

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

	Convey("Given a field and a BSON document, on calling getUpsertValue", t, func() {
		Convey("the value of the key should be correct for unnested documents", func() {
			bsonDocument := bson.D{{"a", 3}}
			So(getUpsertValue("a", bsonDocument), ShouldEqual, 3)
		})
		Convey("the value of the key should be correct for nested document fields", func() {
			inner := bson.D{{"b", 4}}
			bsonDocument := bson.D{{"a", inner}}
			So(getUpsertValue("a.b", bsonDocument), ShouldEqual, 4)
		})
		Convey("the value of the key should be correct for nested document pointer fields", func() {
			inner := bson.D{{"b", 4}}
			bsonDocument := bson.D{{"a", &inner}}
			So(getUpsertValue("a.b", bsonDocument), ShouldEqual, 4)
		})
		Convey("the value of the key should be nil for unnested document "+
			"fields that do not exist", func() {
			bsonDocument := bson.D{{"a", 4}}
			So(getUpsertValue("c", bsonDocument), ShouldBeNil)
		})
		Convey("the value of the key should be nil for nested document "+
			"fields that do not exist", func() {
			inner := bson.D{{"b", 4}}
			bsonDocument := bson.D{{"a", inner}}
			So(getUpsertValue("a.c", bsonDocument), ShouldBeNil)
		})
		Convey("the value of the key should be nil for nested document pointer "+
			"fields that do not exist", func() {
			inner := bson.D{{"b", 4}}
			bsonDocument := bson.D{{"a", &inner}}
			So(getUpsertValue("a.c", bsonDocument), ShouldBeNil)
		})
		Convey("the value of the key should be nil for nil document values", func() {
			So(getUpsertValue("a", bson.D{{"a", nil}}), ShouldBeNil)
		})
	})
}

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

	Convey("Given a set of upsert fields and a BSON document, on calling "+
		"constructUpsertDocument", t, func() {
		Convey("the key/value combination in the upsert document should be "+
			"correct for unnested documents with single fields", func() {
			bsonDocument := bson.D{{"a", 3}}
			upsertFields := []string{"a"}
			upsertDocument := constructUpsertDocument(upsertFields,
				bsonDocument)
			So(upsertDocument, ShouldResemble, bsonDocument)
		})
		Convey("the key/value combination in the upsert document should be "+
			"correct for unnested documents with several fields", func() {
			bsonDocument := bson.D{{"a", 3}, {"b", "string value"}}
			upsertFields := []string{"a"}
			expectedDocument := bson.D{{"a", 3}}
			upsertDocument := constructUpsertDocument(upsertFields,
				bsonDocument)
			So(upsertDocument, ShouldResemble, expectedDocument)
		})
		Convey("the key/value combination in the upsert document should be "+
			"correct for nested documents with several fields", func() {
			inner := bson.D{{testCollection, 4}}
			bsonDocument := bson.D{{"a", inner}, {"b", "string value"}}
			upsertFields := []string{"a.c"}
			expectedDocument := bson.D{{"a.c", 4}}
			upsertDocument := constructUpsertDocument(upsertFields,
				bsonDocument)
			So(upsertDocument, ShouldResemble, expectedDocument)
		})
		Convey("the upsert document should be nil if the key does not exist "+
			"in the BSON document", func() {
			bsonDocument := bson.D{{"a", 3}, {"b", "string value"}}
			upsertFields := []string{testCollection}
			upsertDocument := constructUpsertDocument(upsertFields, bsonDocument)
			So(upsertDocument, ShouldBeNil)
		})
	})
}

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

	Convey("Given a field, its value, and an existing BSON document...", t, func() {
		b := bson.D{{"c", "d"}}
		currentDocument := bson.D{
			{"a", 3},
			{"b", &b},
		}
		Convey("ensure top level fields are set and others, unchanged", func() {
			testDocument := &currentDocument
			expectedDocument := bson.E{"c", 4}
			err := setNestedDocumentValue([]string{"c"}, 4, testDocument, false)
			So(err, ShouldBeNil)
			newDocument := *testDocument
			So(len(newDocument), ShouldEqual, 3)
			So(newDocument[2], ShouldResemble, expectedDocument)
		})
		Convey("ensure new nested top-level fields are set and others, unchanged", func() {
			testDocument := &currentDocument
			expectedDocument := bson.D{{"b", "4"}}
			err := setNestedDocumentValue([]string{"c", "b"}, "4", testDocument, false)
			So(err, ShouldBeNil)
			newDocument := *testDocument
			So(len(newDocument), ShouldEqual, 3)
			So(newDocument[2].Key, ShouldResemble, "c")

			valMap, ok := newDocument[2].Value.(*bson.D)
			So(ok, ShouldBeTrue)

			So(*valMap, ShouldResemble, expectedDocument)
		})
		Convey("ensure existing nested level fields are set and others, unchanged", func() {
			testDocument := &currentDocument
			expectedDocument := bson.D{{"c", "d"}, {"d", 9}}
			err := setNestedDocumentValue([]string{"b", "d"}, 9, testDocument, false)
			So(err, ShouldBeNil)
			newDocument := *testDocument
			So(len(newDocument), ShouldEqual, 2)
			So(newDocument[1].Key, ShouldResemble, "b")

			valMap, ok := newDocument[1].Value.(*bson.D)
			So(ok, ShouldBeTrue)

			So(*valMap, ShouldResemble, expectedDocument)
		})
		Convey("ensure subsequent calls update fields accordingly", func() {
			testDocument := &currentDocument
			expectedDocumentOne := bson.D{{"c", "d"}, {"d", 9}}
			expectedDocumentTwo := bson.E{"f", 23}
			err := setNestedDocumentValue([]string{"b", "d"}, 9, testDocument, false)
			So(err, ShouldBeNil)
			newDocument := *testDocument
			So(len(newDocument), ShouldEqual, 2)
			So(newDocument[1].Key, ShouldResemble, "b")

			valDoc, ok := newDocument[1].Value.(*bson.D)
			So(ok, ShouldBeTrue)

			So(*valDoc, ShouldResemble, expectedDocumentOne)
			err = setNestedDocumentValue([]string{"f"}, 23, testDocument, false)
			So(err, ShouldBeNil)
			newDocument = *testDocument
			So(len(newDocument), ShouldEqual, 3)
			So(newDocument[2], ShouldResemble, expectedDocumentTwo)
		})
	})
}

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

	Convey("Given an unordered BSON document", t, func() {
		Convey("the same document should be returned if there are no blanks", func() {
			bsonDocument := bson.D{{"a", 3}, {"b", "hello"}}
			So(removeBlankFields(bsonDocument), ShouldResemble, bsonDocument)
		})
		Convey("a new document without blanks should be returned if there are "+
			" blanks", func() {
			d := bson.D{
				{"a", ""},
				{"b", ""},
			}
			e := bson.D{
				{"a", ""},
				{"b", 1},
			}
			bsonDocument := bson.D{
				{"a", 0},
				{"b", ""},
				{"c", ""},
				{"d", &d},
				{"e", &e},
			}
			inner := bson.D{
				{"b", 1},
			}
			expectedDocument := bson.D{
				{"a", 0},
				{"e", inner},
			}
			So(removeBlankFields(bsonDocument), ShouldResemble, expectedDocument)
		})
	})
}

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

	Convey("Given an slice of column specs and tokens to convert to BSON", t, func() {
		Convey("the expected ordered BSON should be produced for the given"+
			"column specs and tokens", func() {
			colSpecs := []ColumnSpec{
				{"a", new(FieldAutoParser), pgAutoCast, "auto", []string{"a"}},
				{"b", new(FieldAutoParser), pgAutoCast, "auto", []string{"b"}},
				{"c", new(FieldAutoParser), pgAutoCast, "auto", []string{"c"}},
			}
			tokens := []string{"1", "2", "hello"}
			expectedDocument := bson.D{
				{"a", int32(1)},
				{"b", int32(2)},
				{"c", "hello"},
			}
			bsonD, err := tokensToBSON(colSpecs, tokens, uint64(0), false, false)
			So(err, ShouldBeNil)
			So(bsonD, ShouldResemble, expectedDocument)
		})
		Convey("if there are more tokens than fields, additional fields should be prefixed"+
			" with 'fields' and an index indicating the header number", func() {
			colSpecs := []ColumnSpec{
				{"a", new(FieldAutoParser), pgAutoCast, "auto", []string{"a"}},
				{"b", new(FieldAutoParser), pgAutoCast, "auto", []string{"b"}},
				{"c", new(FieldAutoParser), pgAutoCast, "auto", []string{"c"}},
			}
			tokens := []string{"1", "2", "hello", "mongodb", "user"}
			expectedDocument := bson.D{
				{"a", int32(1)},
				{"b", int32(2)},
				{"c", "hello"},
				{"field3", "mongodb"},
				{"field4", "user"},
			}
			bsonD, err := tokensToBSON(colSpecs, tokens, uint64(0), false, false)
			So(err, ShouldBeNil)
			So(bsonD, ShouldResemble, expectedDocument)
		})
		Convey("an error should be thrown if duplicate headers are found", func() {
			colSpecs := []ColumnSpec{
				{"a", new(FieldAutoParser), pgAutoCast, "auto", []string{"a"}},
				{"b", new(FieldAutoParser), pgAutoCast, "auto", []string{"b"}},
				{"field3", new(FieldAutoParser), pgAutoCast, "auto", []string{"field3"}},
			}
			tokens := []string{"1", "2", "hello", "mongodb", "user"}
			_, err := tokensToBSON(colSpecs, tokens, uint64(0), false, false)
			So(err, ShouldNotBeNil)
		})
		Convey("fields with nested values should be set appropriately", func() {
			colSpecs := []ColumnSpec{
				{"a", new(FieldAutoParser), pgAutoCast, "auto", []string{"a"}},
				{"b", new(FieldAutoParser), pgAutoCast, "auto", []string{"b"}},
				{"c.a", new(FieldAutoParser), pgAutoCast, "auto", []string{"c", "a"}},
			}
			tokens := []string{"1", "2", "hello"}
			c := bson.D{
				{"a", "hello"},
			}
			expectedDocument := bson.D{
				{"a", int32(1)},
				{"b", int32(2)},
				{"c", c},
			}
			bsonD, err := tokensToBSON(colSpecs, tokens, uint64(0), false, false)
			So(err, ShouldBeNil)
			So(expectedDocument[0].Key, ShouldResemble, bsonD[0].Key)
			So(expectedDocument[0].Value, ShouldResemble, bsonD[0].Value)
			So(expectedDocument[1].Key, ShouldResemble, bsonD[1].Key)
			So(expectedDocument[1].Value, ShouldResemble, bsonD[1].Value)
			So(expectedDocument[2].Key, ShouldResemble, bsonD[2].Key)

			valueD, ok := bsonD[2].Value.(*bson.D)
			So(ok, ShouldBeTrue)

			So(expectedDocument[2].Value, ShouldResemble, *valueD)
		})
	})
}

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

	Convey("Given an import worker", t, func() {
		index := uint64(0)
		csvConverters := []CSVConverter{
			{
				colSpecs: []ColumnSpec{
					{"field1", new(FieldAutoParser), pgAutoCast, "auto", []string{"field1"}},
					{"field2", new(FieldAutoParser), pgAutoCast, "auto", []string{"field2"}},
					{"field3", new(FieldAutoParser), pgAutoCast, "auto", []string{"field3"}},
				},
				data:  []string{"a", "b", "c"},
				index: index,
			},
			{
				colSpecs: []ColumnSpec{
					{"field4", new(FieldAutoParser), pgAutoCast, "auto", []string{"field4"}},
					{"field5", new(FieldAutoParser), pgAutoCast, "auto", []string{"field5"}},
					{"field6", new(FieldAutoParser), pgAutoCast, "auto", []string{"field6"}},
				},
				data:  []string{"d", "e", "f"},
				index: index,
			},
		}
		expectedDocuments := []bson.D{
			{
				{"field1", "a"},
				{"field2", "b"},
				{"field3", "c"},
			}, {
				{"field4", "d"},
				{"field5", "e"},
				{"field6", "f"},
			},
		}
		Convey("processDocuments should execute the expected conversion for documents, "+
			"pass then on the output channel, and close the input channel if ordered is true", func() {
			docsInChan := make(chan Converter, 100)
			streamOutChan := make(chan bson.D, 100)
			iw := &importWorker{
				unprocessedDataChan:   docsInChan,
				processedDocumentChan: streamOutChan,
			}
			docsInChan <- csvConverters[0]
			docsInChan <- csvConverters[1]
			close(docsInChan)
			So(iw.processDocuments(t.Context(), true), ShouldBeNil)
			doc1, open := <-streamOutChan
			So(doc1, ShouldResemble, expectedDocuments[0])
			So(open, ShouldEqual, true)
			doc2, open := <-streamOutChan
			So(doc2, ShouldResemble, expectedDocuments[1])
			So(open, ShouldEqual, true)
			_, open = <-streamOutChan
			So(open, ShouldEqual, false)
		})
		Convey("processDocuments should execute the expected conversion for documents, "+
			"pass then on the output channel, and leave the input channel open if ordered is false", func() {
			docsInChan := make(chan Converter, 100)
			streamOutChan := make(chan bson.D, 100)
			iw := &importWorker{
				unprocessedDataChan:   docsInChan,
				processedDocumentChan: streamOutChan,
			}
			docsInChan <- csvConverters[0]
			docsInChan <- csvConverters[1]
			close(docsInChan)
			So(iw.processDocuments(t.Context(), false), ShouldBeNil)
			doc1, open := <-streamOutChan
			So(doc1, ShouldResemble, expectedDocuments[0])
			So(open, ShouldEqual, true)
			doc2, open := <-streamOutChan
			So(doc2, ShouldResemble, expectedDocuments[1])
			So(open, ShouldEqual, true)
			// close will throw a runtime error if streamOutChan is already closed
			close(streamOutChan)
		})
	})
}

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

	Convey(
		"Given some import workers, a Converters input channel and an bson.D output channel",
		t,
		func() {
			docsInChan := make(chan Converter, 5)
			streamOutChan := make(chan bson.D, 5)
			workerInputChannel := []chan Converter{
				make(chan Converter),
				make(chan Converter),
			}
			workerOutputChannel := []chan bson.D{
				make(chan bson.D),
				make(chan bson.D),
			}
			importWorkers := []*importWorker{
				{
					unprocessedDataChan:   workerInputChannel[0],
					processedDocumentChan: workerOutputChannel[0],
				},
				{
					unprocessedDataChan:   workerInputChannel[1],
					processedDocumentChan: workerOutputChannel[1],
				},
			}
			Convey(
				"documents moving through the input channel should be processed and returned in sequence",
				func() {
					eg, ctx := errgroup.WithContext(t.Context())

					// start goroutines to do sequential processing
					for _, iw := range importWorkers {
						eg.Go(
							func() error { return iw.processDocuments(ctx, true) },
						)
					}
					// feed in a bunch of documents
					for _, inputCSVDocument := range csvConverters {
						docsInChan <- inputCSVDocument
					}
					close(docsInChan)
					doSequentialStreaming(ctx, eg, importWorkers, docsInChan, streamOutChan)
					for _, document := range expectedDocuments {
						So(<-streamOutChan, ShouldResemble, document)
					}
				},
			)
		},
	)
}

func TestStreamDocuments(t *testing.T) {
	testtype.SkipUnlessTestType(t, testtype.UnitTestType)
	Convey(`Given:
			1. a boolean indicating streaming order
			2. an input channel where documents are streamed in
			3. an output channel where processed documents are streamed out`, t, func() {

		docsInChan := make(chan Converter, 5)
		streamOutChan := make(chan bson.D, 5)

		Convey(
			"the entire pipeline should complete without error under normal circumstances",
			func() {
				// stream in some documents
				for _, csvConverter := range csvConverters {
					docsInChan <- csvConverter
				}
				close(docsInChan)
				So(streamDocuments(t.Context(), true, 3, docsInChan, streamOutChan), ShouldBeNil)

				// ensure documents are streamed out and processed in the correct manner
				for _, expectedDocument := range expectedDocuments {
					So(<-streamOutChan, ShouldResemble, expectedDocument)
				}
			},
		)
		Convey("the entire pipeline should complete with error if an error is encountered", func() {
			// stream in some documents - create duplicate headers to simulate an error
			csvConverter := CSVConverter{
				colSpecs: []ColumnSpec{
					{"field1", new(FieldAutoParser), pgAutoCast, "auto", []string{"field1"}},
					{"field2", new(FieldAutoParser), pgAutoCast, "auto", []string{"field2"}},
				},
				data:  []string{"a", "b", "c"},
				index: uint64(0),
			}
			docsInChan <- csvConverter
			close(docsInChan)

			// ensure that an error is returned on the error channel
			So(streamDocuments(t.Context(), true, 3, docsInChan, streamOutChan), ShouldNotBeNil)
		})
	})
}
