package tidbautoscaling

import (
	"errors"
	"fmt"
	"time"

	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	infra "github.com/tidbcloud/tidb-management-service/api/serverless_infra"
	infraclient "github.com/tidbcloud/tidb-management-service/api/serverless_infra/client"
	"github.com/tidbcloud/tidb-management-service/api/serverless_infra/meta"
	apitypes "k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

func (ta *TidbAutoscaling) getInfraObjWithToken(token string) (any, error) {
	clusterKey := &meta.ObjectKey{}
	err := ta.infraClusterCli.GetIndex(ta.ctx, token, "status.keyspace_token-index", clusterKey)
	if err == nil {
		cluster := &infra.Cluster{ObjectMeta: meta.ObjectMeta{Namespace: clusterKey.Namespace, Name: clusterKey.Name}}
		if err = ta.infraClusterCli.Get(ta.ctx, cluster); err == nil {
			return cluster, nil
		}
	}
	if !errors.Is(err, infraclient.ErrNotFound) {
		return nil, err
	}

	// for test
	if ta.infraClient == nil {
		return nil, fmt.Errorf("cluster/branch not found")
	}

	branchList := &cloudv1.BranchList{}
	if err := ta.infraClient.List(ta.ctx, branchList, client.MatchingFields(map[string]string{
		"status.fromProvider.keyspaceToken": token,
	})); err != nil {
		return nil, err
	}
	if len(branchList.Items) > 0 {
		return &branchList.Items[0], nil
	}

	return nil, fmt.Errorf("cluster/branch not found")
}

func (ta *TidbAutoscaling) patchTidbReplicas(c any, rg string, replicas int) error {
	switch c := c.(type) {
	case *cloudv1.Branch:
		data := fmt.Sprintf(`{"metadata":{"annotations":{"%s":"%s"}},"spec":{"cluster":{"tidb":{"replicas":%d}}}}`,
			types.AnnoKeyTidbAutoscaling, time.Now().Format(time.RFC3339), replicas)
		return ta.infraClient.Patch(ta.ctx, c, client.RawPatch(apitypes.MergePatchType, []byte(data)))
	case *infra.Cluster:
		set := map[string]any{
			fmt.Sprintf("annotations.'%s'", types.AnnoKeyTidbAutoscaling): time.Now().Format(time.RFC3339),
		}
		if rg == "" {
			set["'spec.tidb'.replicas"] = replicas
		} else if c.Spec.TiDB.ResourceGroups[rg] != nil {
			set[fmt.Sprintf("'spec.tidb'.resource_groups.'%s'.replicas", rg)] = replicas
		} else {
			return fmt.Errorf("resource group %s not found in cluster", rg)
		}
		return ta.infraClusterCli.Patch(ta.ctx, set, c)
	default:
		return fmt.Errorf("type %T is invalid", c)
	}
}