package k8s

import (
	"context"

	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"

	"github.com/tidbcloud/infra-provider/pkg/namer"
	"github.com/tidbcloud/infra-provider/util/collections"
)

// MutClient is a client that can mutate the owner of the object when Create/Update.
type MutClient struct {
	cli   client.Client
	owner client.Object
	namer *namer.Namer
}

func NewMutClientWithOwner(c client.Client, owner client.Object, namer *namer.Namer) *MutClient {
	owner.SetNamespace("")
	return &MutClient{
		cli:   c,
		owner: owner,
	}
}

func (c *MutClient) Create(ctx context.Context, obj client.Object, opts ...client.CreateOption) error {
	if err := controllerutil.SetOwnerReference(c.owner, obj, c.cli.Scheme()); err != nil {
		return err
	}
	obj.SetLabels(collections.MergeMaps(obj.GetLabels(), c.namer.Labels()))
	return c.cli.Create(ctx, obj, opts...)
}

func (c *MutClient) Update(ctx context.Context, obj client.Object, opts ...client.UpdateOption) error {
	if err := controllerutil.SetOwnerReference(c.owner, obj, c.cli.Scheme()); err != nil {
		return err
	}
	obj.SetLabels(collections.MergeMaps(obj.GetLabels(), c.namer.Labels()))
	return c.cli.Update(ctx, obj, opts...)
}

func (c *MutClient) Client() client.Client {
	return c.cli
}