package common

import (
	"context"
	"fmt"
	"maps"

	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	corev1 "k8s.io/api/core/v1"
	storagev1 "k8s.io/api/storage/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
	"github.com/tidbcloud/infra-provider/pkg/util"
)

const gkePersistentDiskProvisioner = "pd.csi.storage.gke.io"

func targetGKEHyperdiskStorageClass(currentStorageClass, desiredStorageClass string) string {
	if !util.GCPHyperdiskPerformanceChanged(currentStorageClass, desiredStorageClass) {
		return ""
	}
	return desiredStorageClass
}

func (t *TiDBCluster) gkeVolumeAttributesClassEnabled() bool {
	return t.Cluster.GetProviderType() == cloudv1.ProviderTypeGKE &&
		t.Cluster.FeatureEnabled(consts.ClusterFeatureUseVolumeAttributesClass) &&
		!t.Cluster.FeatureEnabled(consts.ClusterFeatureUseTiDBOperatorV2)
}

// ensureGKEVolumeAttributesClassesForModification derives VACs from the target
// StorageClasses and commits the copied spec only after every class is validated.
func (t *TiDBCluster) ensureGKEVolumeAttributesClassesForModification(
	ctx context.Context,
	cli client.Client,
	current, desired *opv1alpha1.TidbCluster,
	components []cloudv1.ComponentType,
) error {
	if t.Cluster.GetProviderType() != cloudv1.ProviderTypeGKE {
		return nil
	}

	selectedComponents := sets.New(components...)
	updated := desired.DeepCopy()
	classes := make(map[string]*unstructured.Unstructured)
	prepareClass := func(storageClassName string) error {
		if _, ok := classes[storageClassName]; ok {
			return nil
		}
		storageClass := &storagev1.StorageClass{}
		if err := cli.Get(ctx, client.ObjectKey{Name: storageClassName}, storageClass); err != nil {
			return fmt.Errorf("get target StorageClass %s: %w", storageClassName, err)
		}
		class, err := buildGKEVolumeAttributesClass(storageClass)
		if err != nil {
			return err
		}
		classes[storageClassName] = class
		return nil
	}
	ensureSupported := func() error {
		if t.gkeVolumeAttributesClassEnabled() {
			return nil
		}
		return fmt.Errorf("GKE Hyperdisk performance modification requires VolumeAttributesClass support")
	}

	if selectedComponents.Has(cloudv1.ComponentPD) && current.Spec.PD != nil && desired.Spec.PD != nil {
		targetStorageClass := targetGKEHyperdiskStorageClass(
			ptr.Deref(current.Spec.PD.StorageClassName, ""),
			ptr.Deref(desired.Spec.PD.StorageClassName, ""),
		)
		if targetStorageClass != "" {
			if err := ensureSupported(); err != nil {
				return err
			}
			if err := prepareClass(targetStorageClass); err != nil {
				return err
			}
			updated.Spec.PD.VolumeAttributesClassName = ptr.To(targetStorageClass)
		}
	}

	if selectedComponents.Has(cloudv1.ComponentTiKV) && current.Spec.TiKV != nil && desired.Spec.TiKV != nil {
		currentPrimaryStorageClass := ptr.Deref(current.Spec.TiKV.StorageClassName, "")
		desiredPrimaryStorageClass := ptr.Deref(desired.Spec.TiKV.StorageClassName, "")
		tiKVPrimaryTargetStorageClass := targetGKEHyperdiskStorageClass(
			currentPrimaryStorageClass,
			desiredPrimaryStorageClass,
		)
		additionalVolumeTargets := make(map[string]string)
		if tiKVPrimaryTargetStorageClass == "" {
			currentVolumes := convertVolumeSliceToMap(current.Spec.TiKV.StorageVolumes)
			componentVACActive := ptr.Deref(current.Spec.TiKV.VolumeAttributesClassName, "") != ""
			for i := range desired.Spec.TiKV.StorageVolumes {
				volume := &desired.Spec.TiKV.StorageVolumes[i]
				desiredStorageClass := inheritStorageClass(ptr.Deref(volume.StorageClassName, ""), desiredPrimaryStorageClass)
				currentVolume, exists := currentVolumes[volume.Name]
				if !exists {
					// TiDB Operator defaults a new volume to the component VAC. Pin it to
					// its own effective StorageClass so it cannot inherit the primary disk's attributes.
					if componentVACActive {
						additionalVolumeTargets[volume.Name] = desiredStorageClass
					}
					continue
				}
				currentStorageClass := inheritStorageClass(ptr.Deref(currentVolume.StorageClassName, ""), currentPrimaryStorageClass)
				if targetStorageClass := targetGKEHyperdiskStorageClass(currentStorageClass, desiredStorageClass); targetStorageClass != "" {
					additionalVolumeTargets[volume.Name] = targetStorageClass
				}
			}
		}
		if tiKVPrimaryTargetStorageClass != "" || len(additionalVolumeTargets) > 0 {
			if err := ensureSupported(); err != nil {
				return err
			}

			indexVolumes := func(volumes []opv1alpha1.StorageVolume) (map[string]int, error) {
				indexes := make(map[string]int, len(volumes))
				for i := range volumes {
					name := volumes[i].Name
					if name == "" {
						return nil, fmt.Errorf("TiKV additional volume name must not be empty")
					}
					if _, ok := indexes[name]; ok {
						return nil, fmt.Errorf("duplicate TiKV additional volume %q", name)
					}
					indexes[name] = i
				}
				return indexes, nil
			}
			if _, err := indexVolumes(current.Spec.TiKV.StorageVolumes); err != nil {
				return err
			}
			desiredVolumeIndexes, err := indexVolumes(desired.Spec.TiKV.StorageVolumes)
			if err != nil {
				return err
			}

			if tiKVPrimaryTargetStorageClass != "" {
				for i := range current.Spec.TiKV.StorageVolumes {
					name := current.Spec.TiKV.StorageVolumes[i].Name
					if _, ok := desiredVolumeIndexes[name]; !ok {
						return fmt.Errorf("TiKV storage volume %s is missing from desired TidbCluster", name)
					}
				}
				if err := prepareClass(tiKVPrimaryTargetStorageClass); err != nil {
					return err
				}
				updated.Spec.TiKV.VolumeAttributesClassName = ptr.To(tiKVPrimaryTargetStorageClass)
				// TiDB Operator treats the primary TiKV VAC as a component default,
				// so pin every additional volume to prevent accidental inheritance.
				for i := range desired.Spec.TiKV.StorageVolumes {
					volume := &desired.Spec.TiKV.StorageVolumes[i]
					updatedVolume := &updated.Spec.TiKV.StorageVolumes[i]
					storageClassName := ptr.Deref(volume.StorageClassName, "")
					if storageClassName == "" {
						updatedVolume.VolumeAttributesClassName = ptr.To(tiKVPrimaryTargetStorageClass)
						continue
					}
					if err := prepareClass(storageClassName); err != nil {
						return fmt.Errorf("pin TiKV storage volume %s: %w", volume.Name, err)
					}
					updatedVolume.VolumeAttributesClassName = ptr.To(storageClassName)
				}
			} else {
				for i := range desired.Spec.TiKV.StorageVolumes {
					volume := &desired.Spec.TiKV.StorageVolumes[i]
					targetStorageClass, changed := additionalVolumeTargets[volume.Name]
					if !changed {
						continue
					}
					if err := prepareClass(targetStorageClass); err != nil {
						return fmt.Errorf("pin TiKV storage volume %s: %w", volume.Name, err)
					}
					updated.Spec.TiKV.StorageVolumes[i].VolumeAttributesClassName = ptr.To(targetStorageClass)
				}
			}
		}
	}

	if selectedComponents.Has(cloudv1.ComponentTiFlash) && current.Spec.TiFlash != nil && desired.Spec.TiFlash != nil &&
		len(current.Spec.TiFlash.StorageClaims) > 0 && len(desired.Spec.TiFlash.StorageClaims) > 0 {
		targetStorageClass := targetGKEHyperdiskStorageClass(
			ptr.Deref(current.Spec.TiFlash.StorageClaims[0].StorageClassName, ""),
			ptr.Deref(desired.Spec.TiFlash.StorageClaims[0].StorageClassName, ""),
		)
		if targetStorageClass != "" {
			if err := ensureSupported(); err != nil {
				return err
			}
			if err := prepareClass(targetStorageClass); err != nil {
				return err
			}
			updated.Spec.TiFlash.StorageClaims[0].VolumeAttributesClassName = ptr.To(targetStorageClass)
		}
	}

	if len(classes) == 0 {
		return nil
	}
	for _, class := range classes {
		if err := createOrValidateGKEVolumeAttributesClass(ctx, cli, class); err != nil {
			return err
		}
	}
	desired.Spec = updated.Spec
	return nil
}

func buildGKEVolumeAttributesClass(storageClass *storagev1.StorageClass) (*unstructured.Unstructured, error) {
	iops := storageClass.Parameters["provisioned-iops-on-create"]
	throughput := storageClass.Parameters["provisioned-throughput-on-create"]
	if storageClass.Provisioner != gkePersistentDiskProvisioner ||
		storageClass.Parameters["type"] != consts.GCPVolumeTypeHyperdiskBalanced ||
		iops == "" || throughput == "" {
		return nil, fmt.Errorf("StorageClass %s cannot be converted to a GKE VolumeAttributesClass", storageClass.Name)
	}

	class := &unstructured.Unstructured{Object: map[string]interface{}{
		"driverName": storageClass.Provisioner,
		"parameters": map[string]interface{}{
			"iops":       iops,
			"throughput": throughput,
		},
	}}
	class.SetGroupVersionKind(schema.GroupVersionKind{Group: "storage.k8s.io", Version: "v1", Kind: "VolumeAttributesClass"})
	class.SetName(storageClass.Name)
	return class, nil
}

func createOrValidateGKEVolumeAttributesClass(ctx context.Context, cli client.Client, desiredClass *unstructured.Unstructured) error {
	existingClass := &unstructured.Unstructured{}
	existingClass.SetGroupVersionKind(desiredClass.GroupVersionKind())
	if err := cli.Get(ctx, client.ObjectKeyFromObject(desiredClass), existingClass); err != nil {
		if apierrors.IsNotFound(err) {
			return cli.Create(ctx, desiredClass)
		}
		return err
	}

	desiredDriver, desiredParameters, err := volumeAttributesClassDefinition(desiredClass)
	if err != nil {
		return err
	}
	existingDriver, existingParameters, err := volumeAttributesClassDefinition(existingClass)
	if err != nil {
		return err
	}
	if existingDriver != desiredDriver || !maps.Equal(existingParameters, desiredParameters) {
		return fmt.Errorf("VolumeAttributesClass %s already exists with different driver or parameters", desiredClass.GetName())
	}
	return nil
}

func volumeAttributesClassDefinition(class *unstructured.Unstructured) (string, map[string]string, error) {
	driver, found, err := unstructured.NestedString(class.Object, "driverName")
	if err != nil || !found || driver == "" {
		return "", nil, fmt.Errorf("VolumeAttributesClass %s has invalid driverName", class.GetName())
	}
	parameters, found, err := unstructured.NestedStringMap(class.Object, "parameters")
	if err != nil || !found {
		return "", nil, fmt.Errorf("VolumeAttributesClass %s has invalid parameters", class.GetName())
	}
	return driver, parameters, nil
}

func (t *TiDBCluster) syncGKEVolumeAttributesClassStatus(
	ctx context.Context,
	componentType cloudv1.ComponentType,
	volumeName opv1alpha1.StorageVolumeName,
	desiredStorageClass, desiredClassName string,
	replicas int,
	status *cloudv1.VolumeStatus,
) (bool, error) {
	status.ModifiedStorageClass = desiredStorageClass
	status.ModifiedVolumeAttributesClass = desiredClassName
	if desiredClassName != desiredStorageClass {
		return true, nil
	}

	currentClassName, modifying, err := t.observeVolumeAttributesClass(ctx, componentType, volumeName, desiredClassName, replicas)
	if err != nil {
		return false, err
	}
	status.CurrentVolumeAttributesClass = currentClassName
	if !modifying {
		status.CurrentStorageClass = desiredStorageClass
	}
	return modifying, nil
}

func (t *TiDBCluster) observeVolumeAttributesClass(
	ctx context.Context,
	componentType cloudv1.ComponentType,
	volumeName opv1alpha1.StorageVolumeName,
	desiredClassName string,
	replicas int,
) (string, bool, error) {
	if replicas <= 0 {
		return "", false, fmt.Errorf("invalid expected replicas %d for component %s", replicas, componentType)
	}
	pods, err := k8s.SelectPods(ctx, t.EnvProvider.GetClient(), t.Tidbcluster.Namespace, t.GetComponentPodLabels(componentType))
	if err != nil {
		return "", false, fmt.Errorf("failed to select pods for component %s: %w", componentType, err)
	}
	activePods := make(map[string]struct{})
	for _, pod := range pods {
		if pod.DeletionTimestamp == nil {
			activePods[pod.Name] = struct{}{}
		}
	}

	pvcs, err := k8s.SelectPVCs(ctx, t.EnvProvider.GetClient(), t.Tidbcluster.Namespace, t.GetComponentPodLabels(componentType))
	if err != nil {
		return "", false, fmt.Errorf("failed to select PVCs for component %s: %w", componentType, err)
	}
	var currentClassName string
	matchingPVCs, completedPVCs := 0, 0
	for i := range pvcs {
		pvc := &pvcs[i]
		podName := extractPodNameFromPvcName(pvc.Name, pvc.Labels)
		if _, ok := activePods[podName]; !ok || pvc.Name != fmt.Sprintf("%s-%s", volumeName, podName) {
			continue
		}
		matchingPVCs++
		pvcCurrentClassName := ptr.Deref(pvc.Status.CurrentVolumeAttributesClassName, "")
		if currentClassName == "" && pvcCurrentClassName != desiredClassName {
			currentClassName = pvcCurrentClassName
		}
		complete, err := isPVCVolumeAttributesClassUpdateComplete(pvc, desiredClassName)
		if err != nil {
			return "", false, err
		}
		if complete {
			completedPVCs++
		}
	}
	if len(activePods) == replicas && matchingPVCs == replicas && completedPVCs == replicas {
		return desiredClassName, false, nil
	}
	return currentClassName, true, nil
}

func isPVCVolumeAttributesClassUpdateComplete(pvc *corev1.PersistentVolumeClaim, desiredClassName string) (bool, error) {
	if ptr.Deref(pvc.Spec.VolumeAttributesClassName, "") != desiredClassName {
		return false, nil
	}
	if status := pvc.Status.ModifyVolumeStatus; status != nil {
		if status.TargetVolumeAttributesClassName != desiredClassName {
			return false, nil
		}
		switch status.Status {
		case corev1.PersistentVolumeClaimModifyVolumePending, corev1.PersistentVolumeClaimModifyVolumeInProgress:
			return false, nil
		case corev1.PersistentVolumeClaimModifyVolumeInfeasible:
			return false, fmt.Errorf("PVC %s cannot apply VolumeAttributesClass %q", pvc.Name, desiredClassName)
		default:
			return false, fmt.Errorf("PVC %s reports unknown VolumeAttributesClass modify status %q", pvc.Name, status.Status)
		}
	}
	return ptr.Deref(pvc.Status.CurrentVolumeAttributesClassName, "") == desiredClassName, nil
}