// Code created by gotmpl. DO NOT MODIFY.
// source: internal/shared/otlp/otlptrace/otlptracetest/collector.go.tmpl

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

package otlptracetest

import (
	"cmp"
	"slices"

	collectortracepb "go.opentelemetry.io/proto/otlp/collector/trace/v1"
	commonpb "go.opentelemetry.io/proto/otlp/common/v1"
	resourcepb "go.opentelemetry.io/proto/otlp/resource/v1"
	tracepb "go.opentelemetry.io/proto/otlp/trace/v1"
)

// TracesCollector mocks a collector for the end-to-end testing.
type TracesCollector interface {
	Stop() error
	GetResourceSpans() []*tracepb.ResourceSpans
}

// SpansStorage stores the spans. Mock collectors can use it to
// store spans they have received.
type SpansStorage struct {
	rsm       map[string]*tracepb.ResourceSpans
	spanCount int
}

// NewSpansStorage creates a new spans storage.
func NewSpansStorage() SpansStorage {
	return SpansStorage{
		rsm: make(map[string]*tracepb.ResourceSpans),
	}
}

// AddSpans adds spans to the spans storage.
func (s *SpansStorage) AddSpans(request *collectortracepb.ExportTraceServiceRequest) {
	for _, rs := range request.GetResourceSpans() {
		rstr := resourceString(rs.Resource)
		if existingRs, ok := s.rsm[rstr]; !ok {
			s.rsm[rstr] = rs
			// TODO (rghetia): Add support for library Info.
			if len(rs.ScopeSpans) == 0 {
				rs.ScopeSpans = []*tracepb.ScopeSpans{
					{
						Spans: []*tracepb.Span{},
					},
				}
			}
			s.spanCount += len(rs.ScopeSpans[0].Spans)
		} else {
			if len(rs.ScopeSpans) > 0 {
				newSpans := rs.ScopeSpans[0].GetSpans()
				existingRs.ScopeSpans[0].Spans = append(existingRs.ScopeSpans[0].Spans, newSpans...)
				s.spanCount += len(newSpans)
			}
		}
	}
}

// GetSpans returns the stored spans.
func (s *SpansStorage) GetSpans() []*tracepb.Span {
	spans := make([]*tracepb.Span, 0, s.spanCount)
	for _, rs := range s.rsm {
		spans = append(spans, rs.ScopeSpans[0].Spans...)
	}
	return spans
}

// GetResourceSpans returns the stored resource spans.
func (s *SpansStorage) GetResourceSpans() []*tracepb.ResourceSpans {
	rss := make([]*tracepb.ResourceSpans, 0, len(s.rsm))
	for _, rs := range s.rsm {
		rss = append(rss, rs)
	}
	return rss
}

func resourceString(res *resourcepb.Resource) string {
	sAttrs := sortedAttributes(res.GetAttributes())
	rstr := ""
	for _, attr := range sAttrs {
		rstr = rstr + attr.String()
	}
	return rstr
}

func sortedAttributes(attrs []*commonpb.KeyValue) []*commonpb.KeyValue {
	slices.SortFunc(attrs, func(a, b *commonpb.KeyValue) int {
		return cmp.Compare(a.Key, b.Key)
	})
	return attrs
}
