package helm

import (
	"context"

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

type packageWithPodLabels struct {
	Package
	podLabels map[string]string
}

func WithPodLabels(p Package, labels map[string]string) Package {
	if len(labels) == 0 {
		return p
	}

	copiedLabels := make(map[string]string, len(labels))
	for key, value := range labels {
		copiedLabels[key] = value
	}

	return &packageWithPodLabels{
		Package:   p,
		podLabels: copiedLabels,
	}
}

func ApplyWithPodLabels(ctx context.Context, getter k8s.ClientGetter, p Package, labels map[string]string) error {
	return Apply(ctx, getter, WithPodLabels(p, labels))
}

func (p *packageWithPodLabels) ToHelmValues() map[string]interface{} {
	baseValues := p.Package.ToHelmValues()
	values := map[string]interface{}{}
	for key, value := range baseValues {
		values[key] = value
	}

	podLabels := map[string]interface{}{}
	switch raw := values["podLabels"].(type) {
	case map[string]interface{}:
		for key, value := range raw {
			podLabels[key] = value
		}
	case map[string]string:
		for key, value := range raw {
			podLabels[key] = value
		}
	}
	for key, value := range p.podLabels {
		podLabels[key] = value
	}
	values["podLabels"] = podLabels

	return values
}