package common

import (
	"context"
	"fmt"
	"time"

	"github.com/pingcap/tiflow-operator/api/label"
	"go.uber.org/multierr"
	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/tidbcloud/eks-provider/controllers/replay/task"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/deploy"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
)

type replayTask struct {
	*task.BaseTask
}

func (t *replayTask) buildReplayTolerations() []corev1.Toleration {
	if t == nil || t.Context == nil || t.Context.Provider == nil {
		return nil
	}
	return t.Context.Provider.ReplayTolerations()
}

func (t *replayTask) buildReplayNodeAffinity(replayTask cloudv1.ReplayTask) *corev1.Affinity {
	if t == nil || t.Context == nil || t.Context.Provider == nil {
		return nil
	}
	return t.Context.Provider.ReplayNodeAffinity(replayTask)
}

func NewReplay(ctx *task.Context) deploy.Task {
	return &replayTask{
		BaseTask: &task.BaseTask{Context: ctx},
	}
}

func (t *replayTask) Name() string {
	return "Replay"
}

func (t *replayTask) Create(ctx context.Context) deploy.Result {
	// Create StatefulSet for each task in the replay spec
	// Use V2 (service mode) for deployment
	for _, replayTask := range t.Replay.Spec.Tasks {
		if err := t.ensureStatefulSetV2(ctx, replayTask); err != nil {
			return deploy.Error(err)
		}
	}

	// Update status to Creating - will be updated to Running in PostCreate when all pods are ready
	t.Replay.Status.Phase = cloudv1.ReplayCreating
	t.Replay.Status.Message = fmt.Sprintf("Created %d replay tasks, waiting for pods to be ready", len(t.Replay.Spec.Tasks))

	return deploy.Done()
}

func (t *replayTask) PostCreate(ctx context.Context) deploy.Result {
	// Check StatefulSet status to determine if infrastructure is ready
	clusterCli := t.Remote
	namespace := t.GetReplayNameSpace()

	allReady := true
	totalReplicas := int32(0)
	totalReadyReplicas := int32(0)

	for _, replayTask := range t.Replay.Spec.Tasks {
		stsName := task.GenerateStatefulSetName(t.Replay.Name, replayTask.Name)
		sts := &appsv1.StatefulSet{}
		err := clusterCli.Get(ctx, client.ObjectKey{
			Name:      stsName,
			Namespace: namespace,
		}, sts)

		replicas := replayTask.Replicas
		if replicas == 0 {
			replicas = 1
		}
		totalReplicas += replicas

		if err != nil {
			t.Logger.Info("StatefulSet not found or error", "name", stsName, "error", err)
			allReady = false
			continue
		}

		// Check if all replicas are ready
		if sts.Status.ReadyReplicas < replicas {
			allReady = false
		}

		totalReadyReplicas += sts.Status.ReadyReplicas
	}

	// Check if replay should be stopped based on state
	if t.Replay.Spec.State == cloudv1.ReplayStatePause {
		t.Replay.Status.Phase = cloudv1.ReplayStopped
		t.Replay.Status.Message = "Replay paused by user"
	} else if !allReady {
		// If not all pods are ready, keep phase as Creating
		t.Replay.Status.Phase = cloudv1.ReplayCreating
		t.Replay.Status.Message = fmt.Sprintf("Waiting for pods to be ready: %d/%d replicas ready", totalReadyReplicas, totalReplicas)
	} else {
		// All pods are ready, update phase to Running
		t.Replay.Status.Phase = cloudv1.ReplayRunning
		t.Replay.Status.Message = fmt.Sprintf("All %d replicas are ready and running", totalReplicas)
	}

	// If not all pods are ready, requeue
	if !allReady {
		t.Logger.Info("Waiting for all replay pods to be ready",
			"totalReplicas", totalReplicas,
			"readyReplicas", totalReadyReplicas)
		return deploy.Error(fmt.Errorf("waiting for %d/%d replicas to be ready", totalReadyReplicas, totalReplicas)).WithRequeuAfter(10 * time.Second)
	}

	t.Logger.Info("All replay pods are ready",
		"totalReplicas", totalReplicas,
		"readyReplicas", totalReadyReplicas)
	return deploy.Done()
}

func (t *replayTask) Delete(ctx context.Context) deploy.Result {
	clusterCli := t.Remote
	namespace := t.GetReplayNameSpace()

	// Delete all StatefulSets and PVCs
	var mergeErr error
	for _, replayTask := range t.Replay.Spec.Tasks {
		stsName := task.GenerateStatefulSetName(t.Replay.Name, replayTask.Name)

		// Delete StatefulSet
		sts := &appsv1.StatefulSet{
			ObjectMeta: metav1.ObjectMeta{
				Name:      stsName,
				Namespace: namespace,
			},
		}

		if err := client.IgnoreNotFound(clusterCli.Delete(ctx, sts)); err != nil {
			mergeErr = multierr.Append(mergeErr, err)
		} else {
			t.Logger.Info("Deleting replay StatefulSet", "name", stsName)
		}

		// Delete PVCs if CleanupPolicy is Delete (default behavior)
		// StatefulSet's VolumeClaimTemplates create PVCs with pattern: <volume-name>-<statefulset-name>-<ordinal>
		for i := int32(0); i < replayTask.Replicas; i++ {
			pvcName := fmt.Sprintf("data-%s-%d", stsName, i)
			pvc := &corev1.PersistentVolumeClaim{
				ObjectMeta: metav1.ObjectMeta{
					Name:      pvcName,
					Namespace: namespace,
				},
			}

			if err := client.IgnoreNotFound(clusterCli.Delete(ctx, pvc)); err != nil {
				t.Logger.Error(err, "Failed to delete PVC", "name", pvcName)
				mergeErr = multierr.Append(mergeErr, err)
			} else {
				t.Logger.Info("Deleting replay PVC", "name", pvcName)
			}
		}
	}
	return deploy.ErrorOrDone(mergeErr)
}

func (t *replayTask) PostDelete(ctx context.Context) deploy.Result {
	clusterCli := t.Remote
	namespace := t.GetReplayNameSpace()

	// Delete all StatefulSets and PVCs
	var mergeErr error
	for _, replayTask := range t.Replay.Spec.Tasks {
		stsName := task.GenerateStatefulSetName(t.Replay.Name, replayTask.Name)

		// Delete StatefulSet
		sts := &appsv1.StatefulSet{
			ObjectMeta: metav1.ObjectMeta{
				Name:      stsName,
				Namespace: namespace,
			},
		}

		if err := k8s.DeleteObjectCompleted(ctx, clusterCli, sts); err != nil {
			t.Logger.Info("Waiting for StatefulSet deletion", "name", stsName)
			mergeErr = multierr.Append(mergeErr, err)
		} else {
			t.Logger.Info("Deleted replay StatefulSet", "name", stsName)
		}

		// Delete PVCs if CleanupPolicy is Delete (default behavior)
		// StatefulSet's VolumeClaimTemplates create PVCs with pattern: <volume-name>-<statefulset-name>-<ordinal>
		for i := int32(0); i < replayTask.Replicas; i++ {
			pvcName := fmt.Sprintf("data-%s-%d", stsName, i)
			pvc := &corev1.PersistentVolumeClaim{
				ObjectMeta: metav1.ObjectMeta{
					Name:      pvcName,
					Namespace: namespace,
				},
			}
			if err := k8s.DeleteObjectCompleted(ctx, clusterCli, pvc); err != nil {
				t.Logger.Info("Waiting for PVC deletion", "name", pvcName)
				mergeErr = multierr.Append(mergeErr, err)
			} else {
				t.Logger.Info("Deleted replay PVC", "name", pvcName)
			}
		}
	}
	if mergeErr != nil {
		return deploy.Error(mergeErr).WithRequeuAfter(time.Second * 20)
	}
	return deploy.Done()
}

// ensureStatefulSetV2 creates a StatefulSet with replayer running in service mode (listener mode)
// The replayer will expose HTTP API for creating and canceling replay tasks
func (t *replayTask) ensureStatefulSetV2(ctx context.Context, replayTask cloudv1.ReplayTask) error {
	stsName := task.GenerateStatefulSetName(t.Replay.Name, replayTask.Name)
	namespace := t.GetReplayNameSpace()
	clusterCli := t.Remote

	// Build container args for replayer in service mode
	args := []string{
		"--service-mode",
		"--pprof-addr", fmt.Sprintf("0.0.0.0:%d", task.DefaultSvcPort),
	}

	// Always write logs to a file for persistence
	// We'll use a sidecar to tail the log file to stdout for kubectl logs
	logFilePath := "/data/replayer.log"
	if t.Replay.Spec.ReplayConfig.LogFile != "" {
		logFilePath = t.Replay.Spec.ReplayConfig.LogFile
	}
	if customLogFilePath, ok := t.Replay.Spec.ReplayConfig.CustomArguments["log-file"]; ok && customLogFilePath != "" {
		logFilePath = customLogFilePath
	}
	args = append(args, "--log-file", logFilePath)

	// Get IDs from labels
	tenantID := ""
	projectID := ""
	clusterID := ""
	if t.Replay.Labels != nil {
		tenantID = t.Replay.Labels[cloudv1.TenantTagKey]
		projectID = t.Replay.Labels[cloudv1.ProjectTagKey]
		clusterID = t.Replay.Labels[cloudv1.ClusterTagKey]
	}

	// Build labels
	labels := map[string]string{
		label.NameLabelKey:      t.Replay.Name,
		label.ComponentLabelKey: task.ReplayComponent,
		"replay-name":           t.Replay.Name,
		"replay-task":           replayTask.Name,
		cloudv1.TenantTagKey:    tenantID,
		cloudv1.ProjectTagKey:   projectID,
		cloudv1.ClusterTagKey:   clusterID,
		cloudv1.ReplayTagKey:    t.Replay.Name,
		//"usedby":                "customer-tidb",
	}

	// Add NodeLabels from context for pod affinity
	if t.NodeLabels != nil {
		for k, v := range t.NodeLabels {
			labels[k] = v
		}
	}

	nodeAffinity := t.buildReplayNodeAffinity(replayTask)

	storageSize := task.DefaultDiskSize
	if replayTask.Resources != nil && replayTask.Resources.StorageSize != "" {
		storageSize = replayTask.Resources.StorageSize
	}

	tolerations := t.buildReplayTolerations()

	// Set default replicas if not specified
	replicas := replayTask.Replicas
	if replicas == 0 {
		replicas = 1
	}

	// Build StatefulSet
	sts := &appsv1.StatefulSet{
		ObjectMeta: metav1.ObjectMeta{
			Name:      stsName,
			Namespace: namespace,
			Labels:    labels,
		},
		Spec: appsv1.StatefulSetSpec{
			Replicas: ptr.To(replicas),
			Selector: &metav1.LabelSelector{
				MatchLabels: map[string]string{
					label.NameLabelKey:      t.Replay.Name,
					label.ComponentLabelKey: "replay",
					"replay-task":           replayTask.Name,
				},
			},
			ServiceName: t.Replay.Name, // Use the headless service created by Service task
			Template: corev1.PodTemplateSpec{
				ObjectMeta: metav1.ObjectMeta{
					Labels: labels,
				},
				Spec: corev1.PodSpec{
					Affinity:    t.buildPodAffinity(nodeAffinity, replayTask),
					Tolerations: tolerations,
					InitContainers: []corev1.Container{
						{
							Name:    "sysctl-config",
							Image:   "busybox:latest",
							Command: []string{"sh", "-c"},
							Args: []string{
								`sysctl -w net.ipv4.tcp_tw_reuse=1
sysctl -w net.ipv4.ip_local_port_range="1024 65535"`,
							},
							SecurityContext: &corev1.SecurityContext{
								Privileged: ptr.To(true),
							},
						},
					},
					Containers: []corev1.Container{
						{
							Name:    "replayer",
							Image:   t.Replay.Spec.ReplayConfig.Image,
							Command: []string{"/bin/replayer"},
							Args:    args,
							Ports: []corev1.ContainerPort{
								{
									Name:          "pprof",
									ContainerPort: task.DefaultSvcPort,
									Protocol:      corev1.ProtocolTCP,
								},
							},
							VolumeMounts: []corev1.VolumeMount{
								{
									Name:      "data",
									MountPath: "/data",
								},
							},
						},
						{
							// Log tailer sidecar: tail replayer logs to stdout
							// This allows viewing logs via "kubectl logs -c log-tailer"
							Name:    "log-tailer",
							Image:   "busybox:latest",
							Command: []string{"sh", "-c"},
							Args: []string{
								fmt.Sprintf("tail -F %s", logFilePath),
							},
							VolumeMounts: []corev1.VolumeMount{
								{
									Name:      "data",
									MountPath: "/data",
									ReadOnly:  true, // Read-only access is sufficient for tailing
								},
							},
							Resources: corev1.ResourceRequirements{
								Requests: corev1.ResourceList{
									corev1.ResourceCPU:    resource.MustParse("10m"),
									corev1.ResourceMemory: resource.MustParse("20Mi"),
								},
								Limits: corev1.ResourceList{
									corev1.ResourceCPU:    resource.MustParse("50m"),
									corev1.ResourceMemory: resource.MustParse("50Mi"),
								},
							},
						},
					},
					RestartPolicy: corev1.RestartPolicyAlways,
				},
			},
			// PVC template for data storage
			VolumeClaimTemplates: []corev1.PersistentVolumeClaim{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name: "data",
					},
					Spec: corev1.PersistentVolumeClaimSpec{
						AccessModes: []corev1.PersistentVolumeAccessMode{
							corev1.ReadWriteOnce,
						},
						StorageClassName: ptr.To(task.ReplayStorageClass),
						Resources: corev1.VolumeResourceRequirements{
							Requests: corev1.ResourceList{
								corev1.ResourceStorage: resource.MustParse(storageSize),
							},
						},
					},
				},
			},
		},
	}

	// Apply the StatefulSet
	if err := k8s.Apply(ctx, clusterCli, sts); err != nil {
		return fmt.Errorf("failed to apply StatefulSet %s: %w", stsName, err)
	}

	t.Logger.Info("Ensured replay StatefulSet (service mode)", "name", stsName, "namespace", namespace, "replicas", replicas)
	return nil
}

// buildPodAffinity combines node affinity with pod anti-affinity to ensure:
// 1. Pods are scheduled on the correct nodes (via nodeAffinity)
// 2. Only one replay pod per node (via podAntiAffinity)
func (t *replayTask) buildPodAffinity(nodeAffinity *corev1.Affinity, replayTask cloudv1.ReplayTask) *corev1.Affinity {
	affinity := &corev1.Affinity{}

	// Copy node affinity if provided
	if nodeAffinity != nil {
		affinity.NodeAffinity = nodeAffinity.NodeAffinity
	}
	affinity.PodAntiAffinity = &corev1.PodAntiAffinity{
		RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{
			{
				// Ensure no two replay pods of the same Replay CR run on the same node
				// This applies to all tasks in the Replay CR
				LabelSelector: &metav1.LabelSelector{
					MatchLabels: map[string]string{
						label.NameLabelKey:      t.Replay.Name,
						label.ComponentLabelKey: task.ReplayComponent,
					},
				},
				TopologyKey: corev1.LabelHostname, // kubernetes.io/hostname
			},
		},
	}

	return affinity
}