package enrichment

import (
	"context"
	"log"

	"github.com/tidbcloud/infra-provider/regional-server/pkg/metrics/objectstorage/provider"
)

type BucketLabels struct {
	TenantID     string
	SharedPoolID string
}

type Source interface {
	Load(ctx context.Context) (map[string]BucketLabels, error)
}

type Collector struct {
	base   provider.Collector
	source Source
}

func NewCollector(base provider.Collector, source Source) *Collector {
	return &Collector{base: base, source: source}
}

func (c *Collector) Name() string {
	return c.base.Name()
}

func (c *Collector) Refresh(ctx context.Context) (provider.Snapshot, error) {
	snapshot, err := c.base.Refresh(ctx)
	if err != nil {
		return provider.Snapshot{}, err
	}
	labelsByBucket, err := c.source.Load(ctx)
	if err != nil {
		return provider.Snapshot{}, err
	}
	missingBuckets := map[string]struct{}{}
	for i := range snapshot.Samples {
		labels, ok := labelsByBucket[snapshot.Samples[i].Bucket]
		if !ok {
			if _, warned := missingBuckets[snapshot.Samples[i].Bucket]; !warned {
				log.Printf("object storage bucket %q has no Crossplane bucket enrichment; emitting empty tenant/shared pool labels", snapshot.Samples[i].Bucket)
				missingBuckets[snapshot.Samples[i].Bucket] = struct{}{}
			}
			continue
		}
		snapshot.Samples[i].TenantID = labels.TenantID
		snapshot.Samples[i].SharedPoolID = labels.SharedPoolID
	}
	return snapshot, nil
}