package main

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"cloud.google.com/go/bigquery/storage/managedwriter"
	"google.golang.org/api/option"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protodesc"
	"google.golang.org/protobuf/reflect/protoreflect"
	"google.golang.org/protobuf/types/descriptorpb"
	"google.golang.org/protobuf/types/dynamicpb"
)

const (
	projectID   = "project"
	dataset     = "dataset"
	tableName   = "test_table"
	keyfilePath = "keyfile.json"
)

func setupStream(msgDesc *descriptorpb.DescriptorProto) (*managedwriter.ManagedStream, func() error, error) {
	ctx := context.Background()
	client, err := managedwriter.NewClient(ctx, projectID, option.WithCredentialsFile(keyfilePath))
	if err != nil {
		return nil, nil, err
	}

	tableStr := fmt.Sprintf("projects/%s/datasets/%s/tables/%s", projectID, dataset, tableName)

	managedStream, err := client.NewManagedStream(ctx, managedwriter.WithDestinationTable(tableStr),
		managedwriter.WithType(managedwriter.DefaultStream),
		managedwriter.WithSchemaDescriptor(msgDesc))
	if err != nil {
		return nil, client.Close, err
	}

	return managedStream, client.Close, nil
}

func main() {
	descrProto, msgDescr, err := getDescriptor()
	if err != nil {
		panic(err)
	}

	stream, cleanup, err := setupStream(descrProto)
	if err != nil {
		panic(err)
	}

	defer cleanup()

	nWriters := 5
	nMsg := 10000

	wg := sync.WaitGroup{}
	wg.Add(nWriters)

	for i := 0; i < nWriters; i++ {
		go func() {
			for j := 0; j < nMsg; j++ {
				msg, _ := getMessage(msgDescr, int64(j))

				res, err := stream.AppendRows(context.Background(), [][]byte{msg})
				if err != nil {
					panic(err)
				}

				_, err = res.GetResult(context.Background())
				if err != nil {
					panic(err)
				}

				if j%100 == 0 {
					time.Sleep(5 * time.Second)
				}
			}
			wg.Done()
		}()
	}

	wg.Wait()
}

func getDescriptor() (*descriptorpb.DescriptorProto, protoreflect.MessageDescriptor, error) {
	descrProto := &descriptorpb.DescriptorProto{
		Name: p("record"),
		Field: []*descriptorpb.FieldDescriptorProto{
			{
				Name:     p("c"),
				Number:   p(int32(1)),
				Label:    p(descriptorpb.FieldDescriptorProto_LABEL_OPTIONAL),
				Type:     p(descriptorpb.FieldDescriptorProto_TYPE_INT64),
				JsonName: p("c"),
			},
		},
	}

	fds := &descriptorpb.FileDescriptorSet{
		File: []*descriptorpb.FileDescriptorProto{
			{
				Name:        p("protos/record.proto"),
				Package:     p("main"),
				MessageType: []*descriptorpb.DescriptorProto{descrProto},
				Syntax:      p("proto3"),
			},
		},
	}

	files, err := protodesc.NewFiles(fds)
	if err != nil {
		return nil, nil, err
	}

	desc, err := files.FindDescriptorByName(protoreflect.FullName("main.record"))
	if err != nil {
		return nil, nil, err
	}

	msgDesc, ok := desc.(protoreflect.MessageDescriptor)
	if !ok {
		return nil, nil, errors.New("unable to assert into MessageDescriptor")
	}

	return descrProto, msgDesc, nil
}

func getMessage(msgDesc protoreflect.MessageDescriptor, c int64) ([]byte, error) {
	rqst := dynamicpb.NewMessage(msgDesc)
	fd := rqst.Descriptor().Fields().ByName(protoreflect.Name("c"))
	rqst.Set(fd, protoreflect.ValueOfInt64(c))

	data, err := proto.Marshal(rqst)
	if err != nil {
		return nil, err
	}

	return data, nil
}

func p[T any](t T) *T {
	p := new(T)
	*p = t
	return p
}