// Copyright 2023 Google LLC
//
// 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
//
//     https://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 managedwriter

import (
	"context"
	"io"
	"testing"

	"cloud.google.com/go/bigquery/storage/apiv1/storagepb"
	"cloud.google.com/go/bigquery/storage/managedwriter/testdata"
	"github.com/google/go-cmp/cmp"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protodesc"
	"google.golang.org/protobuf/testing/protocmp"
	"google.golang.org/protobuf/types/descriptorpb"
)

func TestSendOptimizer(t *testing.T) {

	exampleReq := &storagepb.AppendRowsRequest{
		Rows: &storagepb.AppendRowsRequest_ProtoRows{
			ProtoRows: &storagepb.AppendRowsRequest_ProtoData{
				Rows: &storagepb.ProtoRows{
					SerializedRows: [][]byte{[]byte("row_data")},
				},
			},
		},
	}
	exampleStreamID := "foo"
	exampleTraceID := "trace_id"
	exampleReqFull := proto.Clone(exampleReq).(*storagepb.AppendRowsRequest)
	exampleReqFull.WriteStream = exampleStreamID
	exampleReqFull.TraceId = buildTraceID(&streamSettings{TraceID: exampleTraceID})
	exampleDP := &descriptorpb.DescriptorProto{Name: proto.String("schema")}
	exampleReqFull.GetProtoRows().WriterSchema = &storagepb.ProtoSchema{
		ProtoDescriptor: proto.Clone(exampleDP).(*descriptorpb.DescriptorProto),
	}

	ctx := context.Background()

	var testCases = []struct {
		description string
		optimizer   sendOptimizer
		reqs        []*pendingWrite
		sendResults []error
		wantReqs    []*storagepb.AppendRowsRequest
	}{
		{
			description: "verbose-optimizer",
			optimizer:   &verboseOptimizer{},
			reqs: func() []*pendingWrite {
				dv := newDescriptorVersion(exampleDP)
				return []*pendingWrite{
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
				}
			}(),
			sendResults: []error{
				nil,
				io.EOF,
				io.EOF,
			},
			wantReqs: []*storagepb.AppendRowsRequest{
				proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest),
				proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest),
				proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest),
			},
		},
		{
			description: "simplex no errors",
			optimizer:   &simplexOptimizer{},
			reqs: func() []*pendingWrite {
				dv := newDescriptorVersion(exampleDP)
				return []*pendingWrite{
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
				}
			}(),
			sendResults: []error{
				nil,
				nil,
				nil,
			},
			wantReqs: func() []*storagepb.AppendRowsRequest {
				want := make([]*storagepb.AppendRowsRequest, 3)
				// first has no redactions.
				want[0] = proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest)
				req := proto.Clone(want[0]).(*storagepb.AppendRowsRequest)
				req.GetProtoRows().WriterSchema = nil
				req.TraceId = ""
				req.WriteStream = ""
				// second and third are optimized.
				want[1] = req
				want[2] = req
				return want
			}(),
		},
		{
			description: "simplex w/partial errors",
			optimizer:   &simplexOptimizer{},
			reqs: func() []*pendingWrite {
				dv := newDescriptorVersion(exampleDP)
				return []*pendingWrite{
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
				}
			}(),
			sendResults: []error{
				nil,
				io.EOF,
				nil,
			},
			wantReqs: func() []*storagepb.AppendRowsRequest {
				want := make([]*storagepb.AppendRowsRequest, 3)
				want[0] = proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest)
				req := proto.Clone(want[0]).(*storagepb.AppendRowsRequest)
				req.GetProtoRows().WriterSchema = nil
				req.TraceId = ""
				req.WriteStream = ""
				// second request is optimized
				want[1] = req
				// error causes third request to be full again.
				want[2] = want[0]
				return want
			}(),
		},
		{
			description: "multiplex single all errors",
			optimizer:   &multiplexOptimizer{},
			reqs: func() []*pendingWrite {
				dv := newDescriptorVersion(exampleDP)
				return []*pendingWrite{
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
				}
			}(),
			sendResults: []error{
				io.EOF,
				io.EOF,
				io.EOF,
			},
			wantReqs: []*storagepb.AppendRowsRequest{
				proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest),
				proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest),
				proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest),
			},
		},
		{
			description: "multiplex single no errors",
			optimizer:   &multiplexOptimizer{},
			reqs: func() []*pendingWrite {
				dv := newDescriptorVersion(exampleDP)
				return []*pendingWrite{
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
					newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), dv, exampleStreamID, exampleTraceID),
				}
			}(),
			sendResults: []error{
				nil,
				nil,
				nil,
			},
			wantReqs: func() []*storagepb.AppendRowsRequest {
				want := make([]*storagepb.AppendRowsRequest, 3)
				want[0] = proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest)
				req := proto.Clone(want[0]).(*storagepb.AppendRowsRequest)
				req.GetProtoRows().WriterSchema = nil
				req.TraceId = ""
				want[1] = req
				want[2] = req
				return want
			}(),
		},
		{
			description: "multiplex interleave",
			optimizer:   &multiplexOptimizer{},
			reqs: func() []*pendingWrite {
				dvA := newDescriptorVersion(exampleDP)
				dvB := newDescriptorVersion(protodesc.ToDescriptorProto((&testdata.AllSupportedTypes{}).ProtoReflect().Descriptor()))

				reqA := proto.Clone(exampleReq).(*storagepb.AppendRowsRequest)
				reqA.WriteStream = "alpha"

				reqB := proto.Clone(exampleReq).(*storagepb.AppendRowsRequest)
				reqB.WriteStream = "beta"

				writes := make([]*pendingWrite, 10)
				writes[0] = newPendingWrite(ctx, nil, reqA, dvA, reqA.GetWriteStream(), exampleTraceID)
				writes[1] = newPendingWrite(ctx, nil, reqA, dvA, reqA.GetWriteStream(), exampleTraceID)
				writes[2] = newPendingWrite(ctx, nil, reqB, dvB, reqB.GetWriteStream(), exampleTraceID)
				writes[3] = newPendingWrite(ctx, nil, reqA, dvA, reqA.GetWriteStream(), exampleTraceID)
				writes[4] = newPendingWrite(ctx, nil, reqB, dvB, reqB.GetWriteStream(), exampleTraceID)
				writes[5] = newPendingWrite(ctx, nil, reqB, dvB, reqB.GetWriteStream(), exampleTraceID)
				writes[6] = newPendingWrite(ctx, nil, reqB, dvB, reqB.GetWriteStream(), exampleTraceID)
				writes[7] = newPendingWrite(ctx, nil, reqB, dvB, reqB.GetWriteStream(), exampleTraceID)
				writes[8] = newPendingWrite(ctx, nil, reqA, dvA, reqA.GetWriteStream(), exampleTraceID)
				writes[9] = newPendingWrite(ctx, nil, reqA, dvA, reqA.GetWriteStream(), exampleTraceID)

				return writes
			}(),
			sendResults: []error{
				nil,
				nil,
				nil,
				nil,
				nil,
				io.EOF,
				nil,
				nil,
				nil,
				io.EOF,
			},
			wantReqs: func() []*storagepb.AppendRowsRequest {
				want := make([]*storagepb.AppendRowsRequest, 10)

				wantReqAFull := proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest)
				wantReqAFull.WriteStream = "alpha"

				wantReqANoTrace := proto.Clone(wantReqAFull).(*storagepb.AppendRowsRequest)
				wantReqANoTrace.TraceId = ""

				wantReqAOpt := proto.Clone(wantReqAFull).(*storagepb.AppendRowsRequest)
				wantReqAOpt.GetProtoRows().WriterSchema = nil
				wantReqAOpt.TraceId = ""

				wantReqBFull := proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest)
				wantReqBFull.WriteStream = "beta"
				wantReqBFull.GetProtoRows().GetWriterSchema().ProtoDescriptor = protodesc.ToDescriptorProto((&testdata.AllSupportedTypes{}).ProtoReflect().Descriptor())

				wantReqBNoTrace := proto.Clone(wantReqBFull).(*storagepb.AppendRowsRequest)
				wantReqBNoTrace.TraceId = ""

				wantReqBOpt := proto.Clone(wantReqBFull).(*storagepb.AppendRowsRequest)
				wantReqBOpt.GetProtoRows().WriterSchema = nil
				wantReqBOpt.TraceId = ""

				want[0] = wantReqAFull
				want[1] = wantReqAOpt
				want[2] = wantReqBNoTrace
				want[3] = wantReqANoTrace
				want[4] = wantReqBNoTrace
				want[5] = wantReqBOpt
				want[6] = wantReqBFull
				want[7] = wantReqBOpt
				want[8] = wantReqANoTrace
				want[9] = wantReqAOpt

				return want
			}(),
		},
		{
			description: "multiplex w/evolution",
			optimizer:   &multiplexOptimizer{},
			reqs: func() []*pendingWrite {
				dvOld := newDescriptorVersion(exampleDP)
				dvNew := newDescriptorVersion(&descriptorpb.DescriptorProto{Name: proto.String("new")})

				example := proto.Clone(exampleReq).(*storagepb.AppendRowsRequest)

				writes := make([]*pendingWrite, 4)
				writes[0] = newPendingWrite(ctx, nil, example, dvOld, exampleStreamID, exampleTraceID)
				writes[1] = newPendingWrite(ctx, nil, example, dvOld, exampleStreamID, exampleTraceID)
				writes[2] = newPendingWrite(ctx, nil, example, dvNew, exampleStreamID, exampleTraceID)
				writes[3] = newPendingWrite(ctx, nil, example, dvNew, exampleStreamID, exampleTraceID)

				return writes
			}(),
			sendResults: []error{
				nil,
				nil,
				nil,
				nil,
			},
			wantReqs: func() []*storagepb.AppendRowsRequest {
				want := make([]*storagepb.AppendRowsRequest, 4)

				wantBaseReqFull := proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest)

				wantBaseReqOpt := proto.Clone(wantBaseReqFull).(*storagepb.AppendRowsRequest)
				wantBaseReqOpt.TraceId = ""
				wantBaseReqOpt.GetProtoRows().WriterSchema = nil

				wantEvolved := proto.Clone(wantBaseReqOpt).(*storagepb.AppendRowsRequest)
				wantEvolved.GetProtoRows().WriterSchema = &storagepb.ProtoSchema{
					ProtoDescriptor: &descriptorpb.DescriptorProto{Name: proto.String("new")},
				}

				want[0] = wantBaseReqFull
				want[1] = wantBaseReqOpt
				want[2] = wantEvolved
				want[3] = wantBaseReqOpt
				return want
			}(),
		},
	}

	for _, tc := range testCases {
		testARC := &testAppendRowsClient{}
		testARC.sendF = func(req *storagepb.AppendRowsRequest) error {
			testARC.requests = append(testARC.requests, proto.Clone(req).(*storagepb.AppendRowsRequest))
			respErr := tc.sendResults[0]
			tc.sendResults = tc.sendResults[1:]
			return respErr
		}

		for _, req := range tc.reqs {
			err := tc.optimizer.optimizeSend(testARC, req)
			if err != nil {
				tc.optimizer.signalReset()
			}
		}
		// now, compare.
		for k, wr := range tc.wantReqs {
			if diff := cmp.Diff(testARC.requests[k], wr, protocmp.Transform()); diff != "" {
				t.Errorf("%s (req %d) mismatch: -got, +want:\n%s", tc.description, k, diff)
			}
		}
	}
}

func TestDescriptorVersion_EqVersion(t *testing.T) {

	exampleDV := newDescriptorVersion(&descriptorpb.DescriptorProto{Name: proto.String("foo")})
	copiedExampleDV := &descriptorVersion{versionTime: exampleDV.versionTime, descriptorProto: &descriptorpb.DescriptorProto{Name: proto.String("foo")}, hashVal: exampleDV.hashVal}
	exampleDV2 := &descriptorVersion{versionTime: exampleDV.versionTime, descriptorProto: &descriptorpb.DescriptorProto{Name: proto.String("bar")}, hashVal: exampleDV.hashVal}

	testCases := []struct {
		desc    string
		current *descriptorVersion
		other   *descriptorVersion
		want    bool
	}{
		{
			desc: "both nil",
		},
		{
			desc:  "nil current",
			other: newDescriptorVersion(&descriptorpb.DescriptorProto{}),
			want:  false,
		},
		{
			desc:    "nil other",
			current: newDescriptorVersion(&descriptorpb.DescriptorProto{}),
			want:    false,
		},
		{
			desc:    "mismatched",
			current: newDescriptorVersion(&descriptorpb.DescriptorProto{}),
			other:   newDescriptorVersion(&descriptorpb.DescriptorProto{}),
			want:    false,
		},
		{
			desc:    "equal, same reference",
			current: exampleDV,
			other:   exampleDV,
			want:    true,
		},
		{
			desc:    "equal, different references",
			current: exampleDV,
			other:   copiedExampleDV,
			want:    true,
		},
		{
			desc:    "almost equal aka a collision",
			current: exampleDV,
			other:   exampleDV2,
			want:    true,
		},
	}
	for _, tc := range testCases {
		if got := tc.current.eqVersion(tc.other); got != tc.want {
			t.Errorf("case %q, got %t want %t", tc.desc, got, tc.want)
		}
	}
}
