package ack

import (
	"context"
	"fmt"
	"time"

	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/tidbcloud/infra-provider/regional-server/pkg/config"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers/appbox/task"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/deploy"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/util"
)

// CSI driver resource constants
const (
	csiNamespace                  = "kube-system"
	csiProvisionerDeployment      = "csi-provisioner"
	csiPluginDaemonSet            = "csi-plugin"
	csiProvisionerClusterRole     = "alicloud-csi-provisioner"
	provisionerContainer          = "csi-provisioner"
	diskResizerContainer          = "external-disk-resizer"
	externalProvisionerContainer  = "external-disk-provisioner"
	pluginContainer               = "csi-plugin"
	storageAPIGroup               = "storage.k8s.io"
	volumeAttributesClassResource = "volumeattributesclasses"
)

// CSIDriver modifies alicloud csi driver related resources.
// Because VolumeAttributesClass is not supported in ACK yet,
// we fork the csi driver and build our own image to support it.
// This task will do these modifications:
//  1. Modify the deployment kube-system/csi-provisioner:
//     a. Change the image of the container `csi-provisioner` to our own image which supports VolumeAttributesClass.
//     b. Change the image of the container `external-disk-resizer` to our own image which uses storage/v1 VolumeAttributesClass.
//     c. Change the image of the container `external-disk-provisioner` to our own image which uses storage/v1 VolumeAttributesClass.
//     d. Modify the cluster role `alicloud-csi-provisioner` to add VolumeAttributesClass related permissions.
//  2. Modify the daemonset kube-system/csi-plugin:
//     a. Change the image of the container `csi-plugin` to our own image which supports VolumeAttributesClass.
type CSIDriver struct {
	*task.BaseTask
}

func NewCSIDriver(ctx *task.Context) *CSIDriver {
	return &CSIDriver{
		BaseTask: &task.BaseTask{Context: ctx},
	}
}

func (t *CSIDriver) Name() string {
	return "CSIDriver"
}

func (t *CSIDriver) Create(ctx context.Context) deploy.Result {
	t.BeginAt = util.Pointer(time.Now())

	// If CSIDriver config is not set, skip modification
	if t.Config.APPBoxConfig.CSIDriver == nil {
		t.Logger.Info("CSIDriver config not set, skipping CSI driver modification")
		return deploy.Done()
	}

	csiConfig := t.Config.APPBoxConfig.CSIDriver

	// Get remote Kubernetes client
	cli, err := t.Remote.GetClient(ctx)
	if err != nil {
		return deploy.Panic(err)
	}

	var imagePullSecret *corev1.LocalObjectReference
	if requireACKACRImagePullSecret(csiConfig) {
		secret, err := getACKACRImagePullSecret(ctx, cli)
		if err != nil {
			return deploy.Error(fmt.Errorf("failed to get ACR imagePullSecret: %w", err))
		}
		imagePullSecret = &secret
	}

	// Modify csi-provisioner Deployment
	if err := t.modifyCSIProvisionerDeployment(ctx, cli, csiConfig, imagePullSecret); err != nil {
		return deploy.Error(fmt.Errorf("failed to modify csi-provisioner deployment: %w", err))
	}

	// Modify csi-plugin DaemonSet
	if err := t.modifyCSIPluginDaemonSet(ctx, cli, csiConfig, imagePullSecret); err != nil {
		return deploy.Error(fmt.Errorf("failed to modify csi-plugin daemonset: %w", err))
	}

	// Modify alicloud-csi-provisioner ClusterRole
	if err := t.modifyCSIProvisionerClusterRole(ctx, cli); err != nil {
		return deploy.Error(fmt.Errorf("failed to modify alicloud-csi-provisioner clusterrole: %w", err))
	}

	return deploy.Done()
}

// PostCreate ensures the modified CSI resources are ready after creation.
// It will use deploy.Panic() to fail fast if readiness checks fail.
// Otherwise, PVs created by the original CSI driver do NOT have `volumeAttributesClassName` in spec.
func (t *CSIDriver) PostCreate(ctx context.Context) deploy.Result {
	// If CSIDriver config is not set, skip check (consistent with Create behavior)
	if t.Config.APPBoxConfig.CSIDriver == nil {
		return deploy.Done()
	}

	csiConfig := t.Config.APPBoxConfig.CSIDriver
	cli, err := t.Remote.GetClient(ctx)
	if err != nil {
		return deploy.Panic(err)
	}

	// Check csi-provisioner Deployment readiness if any deployment images are configured
	if csiConfig.ProvisionerImage != "" || csiConfig.DiskResizerImage != "" || csiConfig.ExternalProvisionerImage != "" {
		if err := t.ensureDeploymentSynced(ctx, types.NamespacedName{
			Namespace: csiNamespace,
			Name:      csiProvisionerDeployment,
		}, cli); err != nil {
			return deploy.Panic(fmt.Errorf("failed to ensure csi-provisioner deployment rollout: %w", err))
		}
	}

	// Check csi-plugin DaemonSet readiness if plugin image is configured
	if csiConfig.PluginImage != "" {
		if err := t.ensureDaemonSetSynced(ctx, types.NamespacedName{
			Namespace: csiNamespace,
			Name:      csiPluginDaemonSet,
		}, cli); err != nil {
			return deploy.Panic(fmt.Errorf("failed to ensure csi-plugin daemonset rollout: %w", err))
		}
	}

	return deploy.Done()
}

// ensureDeploymentSynced checks if the deployment has completed rollout and the controller
// has observed the latest generation.
func (t *CSIDriver) ensureDeploymentSynced(ctx context.Context, key types.NamespacedName, cli client.Client) error {
	deployment := &appsv1.Deployment{}
	if err := cli.Get(ctx, key, deployment); err != nil {
		return err
	}

	if deployment.Status.ObservedGeneration < deployment.Generation {
		return fmt.Errorf("deployment %s rollout not observed: generation=%d observedGeneration=%d",
			key.String(), deployment.Generation, deployment.Status.ObservedGeneration)
	}

	desired := int32(1)
	if deployment.Spec.Replicas != nil {
		desired = *deployment.Spec.Replicas
	}

	if deployment.Status.UpdatedReplicas < desired {
		return fmt.Errorf("deployment %s rollout in progress: updatedReplicas=%d desiredReplicas=%d",
			key.String(), deployment.Status.UpdatedReplicas, desired)
	}

	if deployment.Status.AvailableReplicas < desired {
		return fmt.Errorf("deployment %s not fully available: availableReplicas=%d desiredReplicas=%d",
			key.String(), deployment.Status.AvailableReplicas, desired)
	}

	return nil
}

// ensureDaemonSetSynced checks if the daemonset has completed rollout and the controller
// has observed the latest generation.
func (t *CSIDriver) ensureDaemonSetSynced(ctx context.Context, key types.NamespacedName, cli client.Client) error {
	ds := &appsv1.DaemonSet{}
	if err := cli.Get(ctx, key, ds); err != nil {
		return err
	}

	if ds.Status.ObservedGeneration < ds.Generation {
		return fmt.Errorf("daemonset %s rollout not observed: generation=%d observedGeneration=%d",
			key.String(), ds.Generation, ds.Status.ObservedGeneration)
	}

	if ds.Status.UpdatedNumberScheduled < ds.Status.DesiredNumberScheduled {
		return fmt.Errorf("daemonset %s rollout in progress: updatedNumberScheduled=%d desiredNumberScheduled=%d",
			key.String(), ds.Status.UpdatedNumberScheduled, ds.Status.DesiredNumberScheduled)
	}

	if ds.Status.NumberReady < ds.Status.DesiredNumberScheduled {
		return fmt.Errorf("daemonset %s not fully ready: numberReady=%d desiredNumberScheduled=%d",
			key.String(), ds.Status.NumberReady, ds.Status.DesiredNumberScheduled)
	}

	return nil
}

// modifyCSIProvisionerDeployment modifies the csi-provisioner deployment to use custom images
func (t *CSIDriver) modifyCSIProvisionerDeployment(ctx context.Context, cli client.Client, cfg *config.CSIDriverConfig, imagePullSecret *corev1.LocalObjectReference) error {
	// Get the deployment
	deployment := &appsv1.Deployment{}
	if err := cli.Get(ctx, types.NamespacedName{
		Namespace: csiNamespace,
		Name:      csiProvisionerDeployment,
	}, deployment); err != nil {
		if apierrors.IsNotFound(err) {
			return fmt.Errorf("deployment %s/%s not found: %w", csiNamespace, csiProvisionerDeployment, err)
		}
		return err
	}

	// Create a deep copy for comparison
	origin := deployment.DeepCopy()
	modified := false

	// Track which containers need to be updated
	var requiredContainers []string
	if cfg.ProvisionerImage != "" {
		requiredContainers = append(requiredContainers, provisionerContainer)
	}
	if cfg.DiskResizerImage != "" {
		requiredContainers = append(requiredContainers, diskResizerContainer)
	}
	if cfg.ExternalProvisionerImage != "" {
		requiredContainers = append(requiredContainers, externalProvisionerContainer)
	}

	// If no images configured for deployment, skip
	if len(requiredContainers) == 0 {
		t.Logger.Info("No deployment images configured, skipping csi-provisioner deployment modification")
		return nil
	}

	if imagePullSecret == nil {
		return fmt.Errorf("ACR imagePullSecret is required but not provided")
	}

	// Verify required containers exist BEFORE modification
	if err := verifyContainersExist(deployment.Spec.Template.Spec.Containers, requiredContainers...); err != nil {
		return fmt.Errorf("container verification failed for deployment %s/%s: %w",
			csiNamespace, csiProvisionerDeployment, err)
	}

	// Ensure imagePullSecrets is set for pulling custom images.
	if ensureImagePullSecret(&deployment.Spec.Template.Spec.ImagePullSecrets, *imagePullSecret) {
		modified = true
	}

	// Update container images
	for i := range deployment.Spec.Template.Spec.Containers {
		container := &deployment.Spec.Template.Spec.Containers[i]
		switch container.Name {
		case provisionerContainer:
			if cfg.ProvisionerImage != "" && container.Image != cfg.ProvisionerImage {
				t.Logger.Info("Updating csi-provisioner container image",
					"from", container.Image, "to", cfg.ProvisionerImage)
				container.Image = cfg.ProvisionerImage
				modified = true
			}
		case diskResizerContainer:
			if cfg.DiskResizerImage != "" && container.Image != cfg.DiskResizerImage {
				t.Logger.Info("Updating external-disk-resizer container image",
					"from", container.Image, "to", cfg.DiskResizerImage)
				container.Image = cfg.DiskResizerImage
				modified = true
			}
		case externalProvisionerContainer:
			if cfg.ExternalProvisionerImage != "" && container.Image != cfg.ExternalProvisionerImage {
				t.Logger.Info("Updating external-disk-provisioner container image",
					"from", container.Image, "to", cfg.ExternalProvisionerImage)
				container.Image = cfg.ExternalProvisionerImage
				modified = true
			}
		}
	}

	// If no changes needed, skip patch
	if !modified {
		t.Logger.Info("csi-provisioner deployment already up to date")
		return nil
	}

	// Patch the deployment
	if err := cli.Patch(ctx, deployment, client.MergeFrom(origin)); err != nil {
		return fmt.Errorf("failed to patch deployment %s/%s: %w", csiNamespace, csiProvisionerDeployment, err)
	}

	t.Logger.Info("Successfully updated csi-provisioner deployment")
	return nil
}

// modifyCSIPluginDaemonSet modifies the csi-plugin daemonset to use custom image
func (t *CSIDriver) modifyCSIPluginDaemonSet(ctx context.Context, cli client.Client, cfg *config.CSIDriverConfig, imagePullSecret *corev1.LocalObjectReference) error {
	// If plugin image not configured, skip
	if cfg.PluginImage == "" {
		t.Logger.Info("Plugin image not configured, skipping csi-plugin daemonset modification")
		return nil
	}

	if imagePullSecret == nil {
		return fmt.Errorf("ACR imagePullSecret is required but not provided")
	}

	// Get the daemonset
	daemonset := &appsv1.DaemonSet{}
	if err := cli.Get(ctx, types.NamespacedName{
		Namespace: csiNamespace,
		Name:      csiPluginDaemonSet,
	}, daemonset); err != nil {
		if apierrors.IsNotFound(err) {
			return fmt.Errorf("daemonset %s/%s not found: %w", csiNamespace, csiPluginDaemonSet, err)
		}
		return err
	}

	// Verify the container exists BEFORE modification
	if err := verifyContainersExist(daemonset.Spec.Template.Spec.Containers, pluginContainer); err != nil {
		return fmt.Errorf("container verification failed for daemonset %s/%s: %w",
			csiNamespace, csiPluginDaemonSet, err)
	}

	// Create a deep copy for comparison
	origin := daemonset.DeepCopy()
	modified := false

	// Ensure imagePullSecrets is set for pulling custom images.
	if ensureImagePullSecret(&daemonset.Spec.Template.Spec.ImagePullSecrets, *imagePullSecret) {
		modified = true
	}

	// Update container image
	for i := range daemonset.Spec.Template.Spec.Containers {
		container := &daemonset.Spec.Template.Spec.Containers[i]
		if container.Name == pluginContainer {
			if container.Image != cfg.PluginImage {
				t.Logger.Info("Updating csi-plugin container image",
					"from", container.Image, "to", cfg.PluginImage)
				container.Image = cfg.PluginImage
				modified = true
			}
			break
		}
	}

	// If no changes needed, skip patch
	if !modified {
		t.Logger.Info("csi-plugin daemonset already up to date")
		return nil
	}

	// Patch the daemonset
	if err := cli.Patch(ctx, daemonset, client.MergeFrom(origin)); err != nil {
		return fmt.Errorf("failed to patch daemonset %s/%s: %w", csiNamespace, csiPluginDaemonSet, err)
	}

	t.Logger.Info("Successfully updated csi-plugin daemonset")
	return nil
}

// modifyCSIProvisionerClusterRole adds VolumeAttributesClass permissions to the ClusterRole
func (t *CSIDriver) modifyCSIProvisionerClusterRole(ctx context.Context, cli client.Client) error {
	// Get the cluster role
	clusterRole := &rbacv1.ClusterRole{}
	if err := cli.Get(ctx, types.NamespacedName{
		Name: csiProvisionerClusterRole,
	}, clusterRole); err != nil {
		if apierrors.IsNotFound(err) {
			return fmt.Errorf("clusterrole %s not found: %w", csiProvisionerClusterRole, err)
		}
		return err
	}

	// Check if VolumeAttributesClass permissions already exist
	if hasVolumeAttributesClassPermissions(clusterRole) {
		t.Logger.Info("VolumeAttributesClass permissions already exist in clusterrole")
		return nil
	}

	// Create a deep copy for comparison
	origin := clusterRole.DeepCopy()

	// Add new policy rule for VolumeAttributesClass
	newRule := rbacv1.PolicyRule{
		APIGroups: []string{storageAPIGroup},
		Resources: []string{volumeAttributesClassResource},
		Verbs:     []string{"get", "list", "watch"},
	}
	clusterRole.Rules = append(clusterRole.Rules, newRule)

	// Patch the cluster role
	if err := cli.Patch(ctx, clusterRole, client.MergeFrom(origin)); err != nil {
		return fmt.Errorf("failed to patch clusterrole %s: %w", csiProvisionerClusterRole, err)
	}

	t.Logger.Info("Successfully added VolumeAttributesClass permissions to clusterrole")
	return nil
}

// hasVolumeAttributesClassPermissions checks if the ClusterRole already has VolumeAttributesClass permissions
func hasVolumeAttributesClassPermissions(clusterRole *rbacv1.ClusterRole) bool {
	for _, rule := range clusterRole.Rules {
		// Check if this rule applies to storage.k8s.io API group or wildcard
		hasAPIGroup := false
		for _, group := range rule.APIGroups {
			if group == storageAPIGroup || group == "*" {
				hasAPIGroup = true
				break
			}
		}
		if !hasAPIGroup {
			continue
		}

		// Check if this rule includes volumeattributesclasses resource or wildcard
		for _, res := range rule.Resources {
			if res == volumeAttributesClassResource || res == "*" {
				return true
			}
		}
	}
	return false
}

// verifyContainersExist verifies that all specified containers exist in the container list
func verifyContainersExist(containers []corev1.Container, requiredNames ...string) error {
	foundContainers := make(map[string]bool)
	for _, container := range containers {
		foundContainers[container.Name] = true
	}

	var missingContainers []string
	for _, name := range requiredNames {
		if !foundContainers[name] {
			missingContainers = append(missingContainers, name)
		}
	}

	if len(missingContainers) > 0 {
		return fmt.Errorf("required containers not found: %v", missingContainers)
	}
	return nil
}

func requireACKACRImagePullSecret(cfg *config.CSIDriverConfig) bool {
	if cfg == nil {
		return false
	}
	return cfg.ProvisionerImage != "" ||
		cfg.DiskResizerImage != "" ||
		cfg.ExternalProvisionerImage != "" ||
		cfg.PluginImage != ""
}

func ensureImagePullSecret(imagePullSecrets *[]corev1.LocalObjectReference, required corev1.LocalObjectReference) bool {
	for _, s := range *imagePullSecrets {
		if s.Name == required.Name {
			return false
		}
	}
	*imagePullSecrets = append(*imagePullSecrets, required)
	return true
}