// Copyright The OpenTelemetry Authors
// SPDX-License-Identifier: Apache-2.0

package test

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/sdk/trace/tracetest"
	"go.opentelemetry.io/otel/trace"

	"go.opentelemetry.io/contrib/instrumentation/go.mongodb.org/mongo-driver/mongo/otelmongo"
)

type validator func(sdktrace.ReadOnlySpan) bool

func TestDBCrudOperation(t *testing.T) {
	commonValidators := []validator{
		func(s sdktrace.ReadOnlySpan) bool {
			return assert.Equal(t, "test-collection.insert", s.Name(), "expected %s", s.Name())
		},
		func(s sdktrace.ReadOnlySpan) bool {
			return assert.Contains(t, s.Attributes(), attribute.String("db.operation.name", "insert"))
		},
		func(s sdktrace.ReadOnlySpan) bool {
			return assert.Contains(t, s.Attributes(), attribute.String("db.collection.name", "test-collection"))
		},
		func(s sdktrace.ReadOnlySpan) bool {
			return assert.Equal(t, codes.Unset, s.Status().Code)
		},
	}

	tt := []struct {
		title          string
		operation      func(context.Context, *mongo.Database) (any, error)
		mockResponses  []bson.D
		excludeCommand bool
		validators     []validator
	}{
		{
			title: "insert",
			operation: func(ctx context.Context, db *mongo.Database) (any, error) {
				return db.Collection("test-collection").InsertOne(ctx, bson.D{{Key: "test-item", Value: "test-value"}})
			},
			mockResponses:  []bson.D{{{Key: "ok", Value: 1}}},
			excludeCommand: false,
			validators: append(commonValidators, func(s sdktrace.ReadOnlySpan) bool {
				for _, attr := range s.Attributes() {
					if attr.Key == "db.query.text" {
						return assert.Contains(t, attr.Value.AsString(), `"test-item":"test-value"`)
					}
				}
				return false
			}),
		},
		{
			title: "insert",
			operation: func(ctx context.Context, db *mongo.Database) (any, error) {
				return db.Collection("test-collection").InsertOne(ctx, bson.D{{Key: "test-item", Value: "test-value"}})
			},
			mockResponses:  []bson.D{{{Key: "ok", Value: 1}}},
			excludeCommand: true,
			validators: append(commonValidators, func(s sdktrace.ReadOnlySpan) bool {
				for _, attr := range s.Attributes() {
					if attr.Key == "db.query.text" {
						return false
					}
				}
				return true
			}),
		},
	}
	for _, tc := range tt {
		title := tc.title
		if tc.excludeCommand {
			title += "/excludeCommand"
		} else {
			title += "/includeCommand"
		}

		mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
		mt.Run(title, func(mt *mtest.T) {
			sr := tracetest.NewSpanRecorder()
			provider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(sr))

			ctx, cancel := context.WithTimeout(t.Context(), time.Second*3)
			defer cancel()

			ctx, span := provider.Tracer("test").Start(ctx, "mongodb-test")

			addr := "mongodb://localhost:27017/?connect=direct"
			opts := options.Client()
			opts.Monitor = otelmongo.NewMonitor(
				otelmongo.WithTracerProvider(provider),
				otelmongo.WithCommandAttributeDisabled(tc.excludeCommand),
			)
			opts.ApplyURI(addr)

			mt.ResetClient(opts)
			mt.AddMockResponses(tc.mockResponses...)

			_, err := tc.operation(ctx, mt.Client.Database("test-database"))
			if err != nil {
				mt.Error(err)
			}

			span.End()

			spans := sr.Ended()
			if !assert.Len(mt, spans, 2, "expected 2 spans, received %d", len(spans)) {
				mt.FailNow()
			}
			assert.Len(mt, spans, 2)
			assert.Equal(mt, spans[0].SpanContext().TraceID(), spans[1].SpanContext().TraceID())
			assert.Equal(mt, spans[0].Parent().SpanID(), spans[1].SpanContext().SpanID())
			assert.Equal(mt, span.SpanContext().SpanID(), spans[1].SpanContext().SpanID())

			s := spans[0]
			assert.Equal(mt, trace.SpanKindClient, s.SpanKind())
			attrs := s.Attributes()
			assert.Contains(mt, attrs, attribute.String("db.system.name", "mongodb"))
			assert.Contains(mt, attrs, attribute.String("network.peer.address", "<mock_connection>:27017"))
			assert.Contains(mt, attrs, attribute.Int64("network.peer.port", int64(27017)))
			assert.Contains(mt, attrs, attribute.String("network.transport", "tcp"))
			assert.Contains(mt, attrs, attribute.String("db.namespace", "test-database"))
			for _, v := range tc.validators {
				assert.True(mt, v(s))
			}
		})
	}
}

func TestDBCollectionAttribute(t *testing.T) {
	tt := []struct {
		title         string
		operation     func(context.Context, *mongo.Database) (any, error)
		mockResponses []bson.D
		validators    []validator
	}{
		{
			title: "delete",
			operation: func(ctx context.Context, db *mongo.Database) (any, error) {
				return db.Collection("test-collection").DeleteOne(ctx, bson.D{{Key: "test-item"}})
			},
			mockResponses: []bson.D{{{Key: "ok", Value: 1}}},
			validators: []validator{
				func(s sdktrace.ReadOnlySpan) bool {
					return assert.Equal(t, "test-collection.delete", s.Name())
				},
				func(s sdktrace.ReadOnlySpan) bool {
					return assert.Contains(t, s.Attributes(), attribute.String("db.operation.name", "delete"))
				},
				func(s sdktrace.ReadOnlySpan) bool {
					return assert.Contains(t, s.Attributes(), attribute.String("db.collection.name", "test-collection"))
				},
				func(s sdktrace.ReadOnlySpan) bool {
					return assert.Equal(t, codes.Unset, s.Status().Code)
				},
			},
		},
		{
			title: "listCollectionNames",
			operation: func(ctx context.Context, db *mongo.Database) (any, error) {
				return db.ListCollectionNames(ctx, bson.D{})
			},
			mockResponses: []bson.D{
				{
					{Key: "ok", Value: 1},
					{Key: "cursor", Value: bson.D{{Key: "firstBatch", Value: bson.A{}}}},
				},
			},
			validators: []validator{
				func(s sdktrace.ReadOnlySpan) bool {
					return assert.Equal(t, "listCollections", s.Name())
				},
				func(s sdktrace.ReadOnlySpan) bool {
					return assert.Contains(t, s.Attributes(), attribute.String("db.operation.name", "listCollections"))
				},
				func(s sdktrace.ReadOnlySpan) bool {
					return assert.Equal(t, codes.Unset, s.Status().Code)
				},
			},
		},
	}
	for _, tc := range tt {
		mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
		mt.Run(tc.title, func(mt *mtest.T) {
			sr := tracetest.NewSpanRecorder()
			provider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(sr))

			ctx, cancel := context.WithTimeout(t.Context(), time.Second*3)
			defer cancel()

			ctx, span := provider.Tracer("test").Start(ctx, "mongodb-test")

			addr := "mongodb://localhost:27017/?connect=direct"
			opts := options.Client()
			opts.Monitor = otelmongo.NewMonitor(
				otelmongo.WithTracerProvider(provider),
				otelmongo.WithCommandAttributeDisabled(true),
			)
			opts.ApplyURI(addr)

			mt.ResetClient(opts)
			mt.AddMockResponses(tc.mockResponses...)

			_, err := tc.operation(ctx, mt.Client.Database("test-database"))
			if err != nil {
				mt.Error(err)
			}

			span.End()

			spans := sr.Ended()
			if !assert.Len(mt, spans, 2, "expected 2 spans, received %d", len(spans)) {
				mt.FailNow()
			}
			assert.Len(mt, spans, 2)
			assert.Equal(mt, spans[0].SpanContext().TraceID(), spans[1].SpanContext().TraceID())
			assert.Equal(mt, spans[0].Parent().SpanID(), spans[1].SpanContext().SpanID())
			assert.Equal(mt, span.SpanContext().SpanID(), spans[1].SpanContext().SpanID())

			s := spans[0]
			assert.Equal(mt, trace.SpanKindClient, s.SpanKind())
			attrs := s.Attributes()
			assert.Contains(mt, attrs, attribute.String("db.system.name", "mongodb"))
			assert.Contains(mt, attrs, attribute.String("network.peer.address", "<mock_connection>:27017"))
			assert.Contains(mt, attrs, attribute.Int64("network.peer.port", int64(27017)))
			assert.Contains(mt, attrs, attribute.String("network.transport", "tcp"))
			assert.Contains(mt, attrs, attribute.String("db.namespace", "test-database"))
			for _, v := range tc.validators {
				assert.True(mt, v(s))
			}
		})
	}
}

func assertSemconv(mt *mtest.T, attrs []attribute.KeyValue) {
	mt.Helper()

	assert.Contains(mt, attrs, attribute.String("db.system.name", "mongodb"))
	assert.Contains(mt, attrs, attribute.String("network.peer.address", "<mock_connection>:27017"))
	assert.Contains(mt, attrs, attribute.Int64("network.peer.port", int64(27017)))
	assert.Contains(mt, attrs, attribute.String("network.transport", "tcp"))
	assert.Contains(mt, attrs, attribute.String("db.namespace", "test-database"))
}

func TestSemanticConventionOptIn(t *testing.T) {
	tt := []struct {
		name         string
		semconvOptIn string
		assert       func(*mtest.T, []attribute.KeyValue)
	}{
		{
			name:         "default",
			semconvOptIn: "",
			assert:       assertSemconv,
		},
		{
			name:         "database",
			semconvOptIn: "database",
			assert:       assertSemconv,
		},
	}
	for _, tc := range tt {
		mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))

		mt.Run(tc.name, func(mt *mtest.T) {
			mt.Setenv("OTEL_SEMCONV_STABILITY_OPT_IN", tc.semconvOptIn)

			sr := tracetest.NewSpanRecorder()
			provider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(sr))

			ctx, cancel := context.WithTimeout(t.Context(), time.Second*3)
			defer cancel()

			ctx, span := provider.Tracer("test").Start(ctx, "mongodb-test")

			addr := "mongodb://localhost:27017/?connect=direct"
			opts := options.Client()
			opts.Monitor = otelmongo.NewMonitor(
				otelmongo.WithTracerProvider(provider),
				otelmongo.WithCommandAttributeDisabled(true),
			)
			opts.ApplyURI(addr)

			mt.ResetClient(opts)
			mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}})

			db := mt.Client.Database("test-database")

			_, _ = db.Collection("test-collection").InsertOne(ctx, bson.D{{Key: "test-item", Value: "test-value"}})

			span.End()

			spans := sr.Ended()
			if !assert.Len(mt, spans, 2, "expected 2 spans, received %d", len(spans)) {
				mt.FailNow()
			}
			assert.Len(mt, spans, 2)
			assert.Equal(mt, spans[0].SpanContext().TraceID(), spans[1].SpanContext().TraceID())
			assert.Equal(mt, spans[0].Parent().SpanID(), spans[1].SpanContext().SpanID())
			assert.Equal(mt, span.SpanContext().SpanID(), spans[1].SpanContext().SpanID())

			s := spans[0]
			assert.Equal(mt, trace.SpanKindClient, s.SpanKind())

			tc.assert(mt, s.Attributes())
		})
	}
}
