package tioperator

import (
	"context"
	"fmt"
	"strings"

	"github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	"github.com/samber/lo"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	schemav1 "github.com/tidbcloud/cluster-service-ng/pkg/storage/etcd/schema/v1"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	auditLogContainerName = "auditlog"
)

// Server labels keys
const (
	LabelKeyKeyspace = "keyspace"
	LabelKeyCidr     = "cidr"
	LabelKeyInternal = "internal.tidbcloud.com"
)

const (
	LabelValueInternalTrue = "true" // Server labels values for internal

	LabelValueKeyspaceSystem = "SYSTEM" // Server labels values
)

func ListTiDBGroups(ctx context.Context, remoteCli client.Client, clusterId string) ([]*entity.Component, error) {
	tidbGroups, err := listTiDBGroups(ctx, remoteCli, dataplaneNamespace(clusterId))
	if err != nil {
		return nil, err
	}

	tidbGroups = lo.Filter(tidbGroups, func(comp *entity.Component, _ int) bool {
		return comp.GetComponentType() == apiv1.ComponentType_TIDB
	})

	return tidbGroups, nil
}

func GetTiDBGroup(ctx context.Context, remoteCli client.Client, clusterId, groupId string) (*entity.Component, error) {
	tidbGroupCr, err := getTiDBGroupCR(ctx, remoteCli, dataplaneNamespace(clusterId), groupId)
	if err != nil {
		if apierrors.IsNotFound(err) {
			return nil, status.Errorf(codes.NotFound, "tidb group %s not found: %s", groupId, err)
		}
		return nil, status.Errorf(codes.Internal, "failed to get tidb group %s: %s", groupId, err)
	}

	return cvtTiDBComponentToEn(ctx, tidbGroupCr)
}

// func GetTiDBAuditLogConfig(ctx context.Context, remoteCli client.Client, clusterId, groupId string) (*entity.AuditLogConfig, error) {
// 	tidbGroupCr, err := getTiDBGroupCR(ctx, remoteCli, dataplaneNamespace(clusterId), groupId)
// 	if err != nil {
// 		if apierrors.IsNotFound(err) {
// 			return nil, status.Errorf(codes.NotFound, "tidb group %s not found: %s", groupId, err)
// 		}
// 		return nil, status.Errorf(codes.Internal, "failed to get tidb group %s: %s", groupId, err)
// 	}

// 	if tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Annotations != nil && tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Annotations[auditLogConfigAnnotation] != "" {
// 		var auditLogConfig entity.AuditLogConfig
// 		err := yaml.Unmarshal([]byte(tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Annotations[auditLogConfigAnnotation]), &auditLogConfig)
// 		if err != nil {
// 			return nil, status.Errorf(codes.Internal, "failed to unmarshal audit log config for tidb group %s: %s", groupId, err)
// 		}
// 		return &auditLogConfig, nil
// 	}

// 	return nil, status.Errorf(codes.NotFound, "audit log config not found for tidb group %s", groupId)
// }

func UpdateTiDBGroup(ctx context.Context, remoteCli client.Client, clusterId, groupId string, options []UpdateOption) (*entity.Component, error) {
	tidbGroupCr, err := getTiDBGroupCR(ctx, remoteCli, dataplaneNamespace(clusterId), groupId)
	if err != nil {
		if apierrors.IsNotFound(err) {
			return nil, status.Errorf(codes.NotFound, "tidb group %s not found: %s", groupId, err)
		}
		return nil, status.Errorf(codes.Internal, "failed to get tidb group %s: %s", groupId, err)
	}

	needUpdate, err := getUpdateExpectedTiDBGroup(ctx, tidbGroupCr, apiv1.ComponentType_TIDB, options...)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to update tidb group %s: %s", groupId, err)
	}

	if needUpdate {
		err = remoteCli.Update(ctx, tidbGroupCr)
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to update tidb group cr %s: %s", groupId, err)
		}
	}

	return cvtTiDBComponentToEn(ctx, tidbGroupCr)
}

func ListTiDBWorkers(ctx context.Context, remoteCli client.Client, clusterId string) ([]*entity.Component, error) {
	// Implementation for listing TiDB workers
	tidbWorkers, err := listTiDBGroups(ctx, remoteCli, dataplaneNamespace(clusterId))
	if err != nil {
		return nil, err
	}
	tidbWorkers = lo.Filter(tidbWorkers, func(comp *entity.Component, _ int) bool {
		return comp.GetComponentType() == apiv1.ComponentType_TIDB_WORKER
	})

	return tidbWorkers, nil
}

func GetTiDBWorker(ctx context.Context, remoteCli client.Client, clusterId, groupId string) (*entity.Component, error) {
	tidbWorkerCr, err := getTiDBWorkerCR(ctx, remoteCli, dataplaneNamespace(clusterId), groupId)
	if err != nil {
		if apierrors.IsNotFound(err) {
			return nil, status.Errorf(codes.NotFound, "tidb worker %s not found: %s", groupId, err)
		}
		return nil, status.Errorf(codes.Internal, "failed to get tidb worker %s: %s", groupId, err)
	}

	return cvtTiDBComponentToEn(ctx, tidbWorkerCr)
}

func UpdateTiDBWorker(ctx context.Context, remoteCli client.Client, clusterId, groupId string, options []UpdateOption) (*entity.Component, error) {
	tidbWorkerCr, err := getTiDBWorkerCR(ctx, remoteCli, dataplaneNamespace(clusterId), groupId)
	if err != nil {
		if apierrors.IsNotFound(err) {
			return nil, status.Errorf(codes.NotFound, "tidb worker %s not found: %s", groupId, err)
		}
		return nil, status.Errorf(codes.Internal, "failed to get tidb worker %s: %s", groupId, err)
	}

	needUpdate, err := getUpdateExpectedTiDBGroup(ctx, tidbWorkerCr, apiv1.ComponentType_TIDB_WORKER, options...)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to update tidb group %s: %s", groupId, err)
	}
	if needUpdate {
		err = remoteCli.Update(ctx, tidbWorkerCr)
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to update tidb worker cr %s: %s", groupId, err)
		}
	}

	return cvtTiDBComponentToEn(ctx, tidbWorkerCr)
}

func listTiDBGroups(ctx context.Context, remoteClient client.Client, dataplaneNs string) ([]*entity.Component, error) {
	// get all component groups
	var tidbGroups v1alpha1.TiDBGroupList
	if err := remoteClient.List(ctx, &tidbGroups, client.InNamespace(dataplaneNs)); err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list tidb groups: %s", err)
	}

	tidbComponents := make([]*entity.Component, 0, len(tidbGroups.Items))

	for i := range tidbGroups.Items {
		compEn, err := cvtTiDBComponentToEn(ctx, &tidbGroups.Items[i])

		if err != nil {
			tilog.Warn(ctx, "convert tidb group to component", zap.Error(err))
		} else {
			tidbComponents = append(tidbComponents, compEn)
		}
	}
	return tidbComponents, nil
}

func getTiDBGroupCR(ctx context.Context, remoteCli client.Client, dataplaneNs, groupId string) (*v1alpha1.TiDBGroup, error) {
	if !strings.HasPrefix(groupId, componentGroupPrefix) {
		groupId = fmt.Sprintf("%s%s", componentGroupPrefix, groupId)
	}

	// get tidb group
	var tidbGroup v1alpha1.TiDBGroup
	key := types.NamespacedName{Name: groupId, Namespace: dataplaneNs}

	if err := remoteCli.Get(
		ctx, key, &tidbGroup,
	); err != nil {
		return nil, err
	}

	return &tidbGroup, nil
}

func getTiDBWorkerCR(ctx context.Context, remoteCli client.Client, dataplaneNs, groupId string) (*v1alpha1.TiDBGroup, error) {
	if !strings.HasPrefix(groupId, componentGroupPrefix) {
		groupId = fmt.Sprintf("%s%s", componentGroupPrefix, groupId)
	}

	// get tidb group
	var tidbGroup v1alpha1.TiDBGroup
	key := types.NamespacedName{Name: fmt.Sprintf("%s-worker", groupId), Namespace: dataplaneNs}

	if err := remoteCli.Get(
		ctx, key, &tidbGroup,
	); err != nil {
		return nil, err
	}

	return &tidbGroup, nil
}

func cvtTiDBComponentToEn(ctx context.Context, tidbGroup *v1alpha1.TiDBGroup) (*entity.Component, error) {
	compEn := &entity.Component{
		Id: strings.Split(tidbGroup.Name, "-")[1], // db-<Id>-tidb
	}
	setTemplateVersionFromAnnotations(compEn, tidbGroup.Annotations)

	switch tidbGroup.Labels[schemav1.LabelComponentKey] {
	case "tidb":
		compEn.ComponentType = apiv1.ComponentType_TIDB
	case "worker-tidb":
		compEn.ComponentType = apiv1.ComponentType_TIDB_WORKER
	default:
		return nil, fmt.Errorf("unknown component type: %s", tidbGroup.Labels[schemav1.LabelComponentKey])
	}

	var err error
	// find the tidb container
	var tidbContainer corev1.Container
	foundContainer := false
	if tidbGroup.Spec.Template.Spec.Overlay != nil &&
		tidbGroup.Spec.Template.Spec.Overlay.Pod != nil {
		for i := range tidbGroup.Spec.Template.Spec.Overlay.Pod.Spec.Containers {
			if tidbGroup.Spec.Template.Spec.Overlay.Pod.Spec.Containers[i].Name == "tidb" {
				tidbContainer = tidbGroup.Spec.Template.Spec.Overlay.Pod.Spec.Containers[i]
				foundContainer = true
				break
			}
		}
	}
	if !foundContainer {
		return nil, fmt.Errorf("container named 'tidb' not found in TiDBGroup %s", tidbGroup.Name)
	}

	// Get nodeSelector from Overlay if present
	var nodeSelector map[string]string
	if tidbGroup.Spec.Template.Spec.Overlay != nil &&
		tidbGroup.Spec.Template.Spec.Overlay.Pod != nil {
		nodeSelector = tidbGroup.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector
	}

	compEn.NodeQuota, err = CvtNodeSpecToQuota(
		tidbGroup.Labels,
		tidbGroup.Annotations,
		uint32(*tidbGroup.Spec.Replicas),
		tidbContainer.Resources.Limits,
		nodeSelector,
	)
	if err != nil {
		return nil, fmt.Errorf("failed to get tidb group %s quota: %w", compEn.GetId(), err)
	}
	compEn.Version = tidbGroup.Spec.Template.Spec.Version
	imageRepo, imageTag := splitImageRepoTag(tidbGroup.Spec.Template.Spec.Image)
	compEn.Configuration = &entity.ComponentConfiguration{
		ImageRepo:                imageRepo,
		ImageTag:                 imageTag,
		ConfigurationFileContent: string(tidbGroup.Spec.Template.Spec.Config),
	}
	compEn.Storages, err = cvtStoragesToEn(ctx, tidbGroup.Spec.Template.Spec.Volumes)
	if err != nil {
		return nil, err
	}
	compEn.Phase = getTiDBGroupStatus(ctx, tidbGroup)
	// set labels from server labels
	if tidbGroup.Spec.Template.Spec.Server.Labels != nil {
		compEn.Labels = tidbGroup.Spec.Template.Spec.Server.Labels
	}
	return compEn, nil
}

func getTiDBGroupStatus(ctx context.Context, tidbGroup *v1alpha1.TiDBGroup) apiv1.Component_Phase {
	// if synced is false, it means the tidbgroup is not synced yet
	if !isSynced(tidbGroup.Status.Conditions) {
		tilog.Warn(ctx, "tidbgroup is not synced", zap.Any("Conditions", tidbGroup.Status.Conditions))
		return apiv1.Component_MODIFYING
	}

	if !isUpToDate(tidbGroup.Generation, tidbGroup.Status.CommonStatus) {
		tilog.Warn(ctx, "tidbgroup is not up to date", zap.Any("Status", tidbGroup.Status))
		return apiv1.Component_MODIFYING
	}

	if tidbGroup.Status.ReadyReplicas == 0 || tidbGroup.Status.UpdatedReplicas == 0 {
		tilog.Warn(ctx, "tidbGroup has zero ready or updated replicas", zap.Int32("ReadyReplicas", tidbGroup.Status.ReadyReplicas), zap.Int32("UpdatedReplicas", tidbGroup.Status.UpdatedReplicas))
		return apiv1.Component_UNAVAILABLE
	}

	// check if tidbgroup is ready
	// if ready is false, it means the tidbgroup is unavailable
	if !isReady(tidbGroup.Status.Conditions) {
		tilog.Warn(ctx, "tidbgroup is not ready", zap.Any("Conditions", tidbGroup.Status.Conditions))
		return apiv1.Component_UNAVAILABLE
	}

	return apiv1.Component_AVAILABLE
}

func getUpdateExpectedTiDBGroup(ctx context.Context, tidbGroupCr *v1alpha1.TiDBGroup, compType apiv1.ComponentType, options ...UpdateOption) (bool, error) {
	needUpdate := false
	for _, opt := range options {
		switch opt.Option {
		case OpsUpdateNodeCount:
			count, ok := opt.Value.(uint32)
			if !ok {
				tilog.Error(ctx, "invalid node count", zap.Any("value", opt.Value))
				return false, fmt.Errorf("invalid node count for tidb group %s", tidbGroupCr.Name)
			}
			if count > 0 {
				currNodeQuota, err := CvtNodeSpecToQuota(tidbGroupCr.Labels, tidbGroupCr.Annotations, uint32(*tidbGroupCr.Spec.Replicas), nil, nil)
				if err != nil {
					return false, status.Errorf(codes.Internal, "failed to convert node spec to quota for tikv worker %s: %s", tidbGroupCr.Name, err)
				}

				nodeQuota := &entity.NodeQuota{
					NodeCount: count,
				}

				err = checkNodeQuota(currNodeQuota, nodeQuota)
				if err != nil {
					return false, status.Errorf(codes.InvalidArgument, "invalid node quota for tikv worker %s: %s", tidbGroupCr.Name, err)
				}

				if uint32(*tidbGroupCr.Spec.Replicas) != count {
					tidbGroupCr.Spec.Replicas = lo.ToPtr(int32(count))
					needUpdate = true
				}
			}

		case OpsUpdateNodeSpec:
			nodeSpec, ok := opt.Value.(*entity.NodeSpec)
			if !ok {
				tilog.Error(ctx, "invalid node spec", zap.Any("value", opt.Value))
				return false, fmt.Errorf("invalid node spec for tidb group %s", tidbGroupCr.Name)
			}

			resourceRequirements, err := setNodeSpec(nodeSpec, tidbGroupCr.Labels, tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector, compType)
			if err != nil {
				tilog.Error(ctx, "set node spec", zap.Error(err))
				return false, fmt.Errorf("failed to set node spec for tidb group %s: %w", tidbGroupCr.Name, err)
			}
			for i := range tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Spec.Containers {
				if tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Spec.Containers[i].Name != "tidb" {
					continue
				}

				tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Spec.Containers[i].Resources = resourceRequirements

				limitMem := EffectiveMaxMemoryMi(nodeSpec)
				for j := range tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Spec.Containers[i].Env {
					if tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Spec.Containers[i].Env[j].Name == "GOMEMLIMIT" {
						tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Spec.Containers[i].Env[j].Value = fmt.Sprintf(
							"%dMiB", uint64(float64(limitMem)*GoMemlimitRate))
					}
				}
				needUpdate = true
				break
			}

		case OpsUpdateMaxCount:
			maxCount, ok := opt.Value.(uint32)
			if !ok {
				tilog.Error(ctx, "invalid max count", zap.Any("value", opt.Value))
				return false, fmt.Errorf("invalid max count for tidb group %s", tidbGroupCr.Name)
			}
			if maxCount > 0 {
				annotations := tidbGroupCr.Annotations
				if annotations == nil {
					annotations = map[string]string{}
				}

				annotations[schemav1.AnnotationsComponentMaxCount] = fmt.Sprintf("%d", maxCount)
				tidbGroupCr.Annotations = annotations

				needUpdate = true
			}
		case OpsUpdateMinCount:
			minCount, ok := opt.Value.(uint32)
			if !ok {
				tilog.Error(ctx, "invalid min count", zap.Any("value", opt.Value))
			}

			annotations := tidbGroupCr.Annotations
			if annotations == nil {
				annotations = map[string]string{}
			}

			annotations[schemav1.AnnotationsComponentMinCount] = fmt.Sprintf("%d", minCount)
			tidbGroupCr.Annotations = annotations

			needUpdate = true
		case OpsUpdateVersion:
			version, ok := opt.Value.(string)
			if !ok {
				tilog.Error(ctx, "invalid version", zap.Any("value", opt.Value))
				return false, fmt.Errorf("invalid version for tidb group %s", tidbGroupCr.Name)
			}
			if version != "" {
				tidbGroupCr.Spec.Template.Spec.Version = version
				needUpdate = true
			}
		case OpsUpdateImage:
			image, ok := opt.Value.(string)
			if !ok {
				tilog.Error(ctx, "invalid image", zap.Any("value", opt.Value))
				return false, fmt.Errorf("invalid image for tidb group %s", tidbGroupCr.Name)
			}
			if image != "" {
				tidbGroupCr.Spec.Template.Spec.Image = &image
				needUpdate = true
			}
		case OpsUpdateConfig:
			config, ok := opt.Value.(string)
			if !ok {
				tilog.Error(ctx, "invalid config", zap.Any("value", opt.Value))
				return false, fmt.Errorf("invalid config for tidb group %s", tidbGroupCr.Name)
			}
			if config != "" {
				tidbGroupCr.Spec.Template.Spec.Config = v1alpha1.ConfigFile(config)
				needUpdate = true
			}
		case OpsUpdateUpdateStrategy:
			strategy, ok := opt.Value.(v1alpha1.ConfigUpdateStrategy)
			if !ok {
				tilog.Error(ctx, "invalid update strategy", zap.Any("value", opt.Value))
				return false, fmt.Errorf("invalid update strategy for tidb group %s", tidbGroupCr.Name)
			}
			if strategy != "" {
				tidbGroupCr.Spec.Template.Spec.UpdateStrategy.Config = strategy
				needUpdate = true
			}
		case OpsUpdateTiDBAuditLog:
			auditLogConfig, ok := opt.Value.(string)
			if !ok {
				tilog.Error(ctx, "invalid audit log config", zap.Any("value", opt.Value))
				return false, fmt.Errorf("invalid audit log config for tidb group %s", tidbGroupCr.Name)
			}
			if auditLogConfig != "" {
				annotations := tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Annotations
				if annotations == nil {
					annotations = map[string]string{}
				}
				annotations[schemav1.AuditLogConfigAnnotation] = auditLogConfig
				tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Annotations = annotations
				needUpdate = true
			}
		case OpsUpdateTiDBAuditLogImage:
			targetImage, ok := opt.Value.(string)
			if !ok {
				tilog.Error(ctx, "invalid audit log image", zap.Any("value", opt.Value))
				return false, fmt.Errorf("invalid audit log image for tidb group %s", tidbGroupCr.Name)
			}
			if targetImage != "" {
				if tidbGroupCr.Spec.Template.Spec.Overlay == nil ||
					tidbGroupCr.Spec.Template.Spec.Overlay.Pod == nil ||
					tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Spec == nil {
					tilog.Warn(ctx, "pod overlay spec not found, skip audit log image update", zap.String("tidbGroup", tidbGroupCr.Name))
					continue
				}
				containers := tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Spec.Containers
				for i := range containers {
					if containers[i].Name == auditLogContainerName {
						if containers[i].Image != targetImage {
							tilog.Info(ctx, "updating audit log sidecar image",
								zap.String("tidbGroup", tidbGroupCr.Name),
								zap.String("from", containers[i].Image),
								zap.String("to", targetImage),
							)
							containers[i].Image = targetImage
							tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Spec.Containers = containers
							needUpdate = true
						}
						break
					}
				}
			}
		case OpsUpdateLabels:
			labels, ok := opt.Value.(map[string]string)
			if !ok {
				tilog.Error(ctx, "invalid labels", zap.Any("value", opt.Value))
				return false, fmt.Errorf("invalid labels for tidb group %s", tidbGroupCr.Name)
			}
			if labels != nil {
				serverLabels := tidbGroupCr.Spec.Template.Spec.Server.Labels
				if serverLabels == nil {
					serverLabels = map[string]string{}
				}
				// merge labels
				for k, v := range labels {
					serverLabels[k] = v
				}
				tidbGroupCr.Spec.Template.Spec.Server.Labels = serverLabels
				needUpdate = true
			}
		case OpsUpdateLabelCidr:
			cidr, ok := opt.Value.(string)
			if !ok {
				tilog.Error(ctx, "invalid cidr", zap.Any("value", opt.Value))
				return false, fmt.Errorf("invalid cidr for tidb group %s", tidbGroupCr.Name)
			}
			if cidr != "" {
				labels := tidbGroupCr.Spec.Template.Spec.Server.Labels
				if labels == nil {
					labels = map[string]string{}
				}
				labels[LabelKeyCidr] = cidr
				tidbGroupCr.Spec.Template.Spec.Server.Labels = labels
				needUpdate = true
			}
		case OpsUpdateNodeVersion:
			nodeVersion, ok := opt.Value.(string)
			if !ok || nodeVersion == "" {
				tilog.Error(ctx, "invalid node version", zap.Any("value", opt.Value))
				return false, status.Errorf(codes.InvalidArgument, "invalid node version for tidb group %s", tidbGroupCr.Name)
			}
			nodeSelector := tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector
			if nodeSelector == nil {
				nodeSelector = map[string]string{}
			}
			if nodeSelector[schemav1.K8sVersionLabelKey] != nodeVersion {
				nodeSelector[schemav1.K8sVersionLabelKey] = nodeVersion
				tidbGroupCr.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector = nodeSelector
				needUpdate = true
			}
		}
	}

	return needUpdate, nil
}