package export

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"sync/atomic"
	"time"

	"github.com/go-logr/logr"
	"github.com/pingcap/errors"
	"github.com/tidbcloud/aws-shared-provider/pkg/keyspacemigration"
	"github.com/tidbcloud/aws-shared-provider/pkg/metrics"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/pd"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	infraclient "github.com/tidbcloud/tidb-management-service/api/serverless_infra/client"
	clientv3 "go.etcd.io/etcd/client/v3"
	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	k8slabels "k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/serializer"
	apitypes "k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/tools/record"
	"k8s.io/client-go/util/workqueue"
	ctrl "sigs.k8s.io/controller-runtime"
	ctrlcache "sigs.k8s.io/controller-runtime/pkg/cache"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/cluster"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/event"
	"sigs.k8s.io/controller-runtime/pkg/handler"
	ctrlmanager "sigs.k8s.io/controller-runtime/pkg/manager"
	"sigs.k8s.io/controller-runtime/pkg/predicate"
	"sigs.k8s.io/controller-runtime/pkg/source"
)

const (
	MaxConcurrentExport      = 10
	maxConcurrentExportsEnv  = "MAX_CONCURRENT_EXPORTS"
	ExportTiDBPoolNS         = "export-tidb-pool"
	ExportVipTiDBPoolNS      = "export-vip-tidb-pool"
	ExportSuperVipTiDBPoolNS = "export-super-vip-tidb-pool"
	DumplingContainerName    = "dumpling"
	DumplingJobType          = "dumpling"
	requeueAfter             = 15 * time.Second
	ExportAssignTimeout      = 60 * time.Second

	// The export TiDB pools run with min replicas 0, so the first assign only triggers an
	// on-demand pod creation and fails. Re-queue instead of blocking the reconcile worker,
	// until the pod becomes ready or the whole export exceeds DefaultAssignTiDBMaxWait.
	// The budget is generous on purpose: a stuck export stays visible and recoverable by
	// hand for an hour before it is failed for good.
	assignTiDBRequeueInterval = 15 * time.Second
	DefaultAssignTiDBMaxWait  = time.Hour
)

type UnRetryableError struct {
	message string
}

func (e *UnRetryableError) Error() string {
	return fmt.Sprintf("error message=%s", e.message)
}

// Reconciler reconciles a Export object.
type Reconciler struct {
	logr.Logger

	InfraClusterCli infraclient.Client

	Client        client.Client
	InfraClient   client.Client
	WorkloadCache ctrlcache.Cache

	InfraAPIReader client.Reader

	PdClient   pd.HttpPdClient
	HttpClient *http.Client

	Recorder      record.EventRecorder
	InfraRecorder record.EventRecorder

	EtcdKV             clientv3.KV
	KubeClusterName    string
	EtcdTokenPrefixMap map[string]string
	TidbVersionMap     map[string]string
	TidbPools          []string
	ProviderType       string

	DebugMode          bool
	Codecs             serializer.CodecFactory
	ExportJobCM        string
	ExportJobNS        string
	DumplingImage      string
	ExportStashingArea string
	// exportJob is immutable after publication. Reconciles must DeepCopy it
	// before applying export-specific configuration.
	exportJob atomic.Pointer[batchv1.Job]

	AWSAccountID string

	EnableTLS bool

	MigrationLocker keyspacemigration.Locker
}

func (r *Reconciler) initRuntimeWithExportCluster(exportCRCluster cluster.Cluster) {
	r.InfraClient = exportCRCluster.GetClient()
	r.InfraAPIReader = exportCRCluster.GetAPIReader()
	r.InfraRecorder = exportCRCluster.GetEventRecorderFor("Export")
}

// reconcileContext is the local context of a single cluster reconciliation, it is NOT concurrency safe.
type reconcileContext struct {
	*Reconciler
	logr.Logger

	ctx    context.Context
	export *cloudv1.Export

	projectID string
	tenantID  string
	clusterID string

	keyspaceToken string
	keyspaceID    uint32

	poolSelector string
}

func (r *Reconciler) LogDebug(ctx context.Context, msg string) {
	log := ctrl.LoggerFrom(ctx)
	if r.DebugMode {
		log.Info(msg)
	}
}

func (r *Reconciler) newReconcileContext(ctx context.Context, req ctrl.Request) (*reconcileContext, error) {
	export := &cloudv1.Export{}
	if err := r.InfraClient.Get(ctx, req.NamespacedName, export); err != nil {
		if apierrors.IsNotFound(err) {
			return nil, nil
		}
		return nil, fmt.Errorf("can't get export: %w", err)
	}

	if kubeCluster := labels.GetKubeClusterFromLabels(export.Labels); kubeCluster != r.KubeClusterName {
		r.Info("unexpected cluster scheduling", "kubecluster", kubeCluster)
		return nil, nil
	}

	exportMeta := &types.ExportMeta{}
	if err := json.Unmarshal([]byte(export.Annotations[types.AnnoExportMeta]), exportMeta); err != nil {
		return nil, fmt.Errorf("can't unmarshal export meta: %w", err)
	}
	keyspaceID := exportMeta.KeyspaceID
	keyspaceToken := exportMeta.KeyspaceToken

	r.LogDebug(ctx, fmt.Sprintf("keyspaceToken: %s", keyspaceToken))

	return &reconcileContext{
		Reconciler: r,
		Logger:     ctrl.LoggerFrom(ctx),

		ctx:           ctx,
		export:        export,
		keyspaceToken: keyspaceToken,
		keyspaceID:    keyspaceID,
		poolSelector:  exportMeta.TiDBPoolSelector,

		projectID: exportMeta.ProjectID,
		tenantID:  exportMeta.TenantID,
		clusterID: exportMeta.ClusterID,
	}, nil
}

// Reconcile is part of the main kubernetes reconciliation loop which aims to
// move the current state of the export closer to the desired state.
func (r *Reconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	log := ctrl.LoggerFrom(ctx)
	log.Info("start reconcile export")
	start := time.Now()
	defer func() {
		log.Info("finish reconcile export", "time", time.Since(start))
	}()

	rx, err := r.newReconcileContext(ctx, req)
	r.LogDebug(ctx, "new reconcile context success")
	if err != nil || rx == nil {
		return utils.ReconcileErr(log, "Export", req.Name, err)
	}

	// ensure finalizer
	action, err := rx.ensureFinalizer()
	if action || err != nil {
		return utils.ReconcileErr(log, "Export", rx.export.Name, err)
	}

	r.LogDebug(ctx, "ensure finalizer success")

	// shortcut to Cancel phase
	if rx.export.Annotations[types.AnnoExportCancel] != "" && rx.export.Status.Phase != cloudv1.ExportPhaseCancel {
		rx.Info("start to cancel export")
		action, err = rx.updateProgress(cloudv1.ExportPhaseCancel, "export is canceled")
		if action || err != nil {
			return utils.ReconcileErr(log, "Export", rx.export.Name, err)
		}
	}

	// shortcut to Pending phase when export is newly created
	if rx.export.Status.Phase == "" {
		rx.Info("start to update phase to prepare")
		action, err = rx.updateProgress(cloudv1.ExportPhasePrepare, "export is waiting to be reconciled")
		if action || err != nil {
			return utils.ReconcileErr(log, "Export", rx.export.Name, err)
		}
	}

	if rx.export.Annotations[types.AnnoExportMetering] != "" &&
		rx.export.Annotations[types.AnnoExportConsoleMetric] == "" {
		// get usage from annotation
		usage, err := strconv.ParseFloat(rx.export.Annotations[types.AnnoExportMetering], 64)
		if err != nil {
			rx.Error(err, "failed to parse float for usage", "usage", rx.export.Annotations[types.AnnoExportMetering])
			usage = -1
		}
		// report console metric
		if usage > 0 && rx.keyspaceID != 0 {
			metrics.ExportRUGauge.WithLabelValues(rx.export.Spec.ClusterName, rx.export.Name,
				strconv.Itoa(int(rx.keyspaceID))).Set(usage)
		}
		// update cr
		rx.export.Annotations[types.AnnoExportConsoleMetric] = rx.export.Annotations[types.AnnoExportMetering]
		err = r.InfraClient.Update(rx.ctx, rx.export)
		if err != nil {
			rx.Error(err, "failed to update export annotation")
			return utils.ReconcileErr(log, "Export", rx.export.Name, err)
		}
		return ctrl.Result{}, nil
	}

	// start reconcile
	exportPhase := rx.export.Status.Phase
	switch exportPhase {
	case cloudv1.ExportPhasePrepare:
		// prepare export
		err = rx.startExport()
	case cloudv1.ExportPhaseRunning:
		// export is running, watch dumpling job
		err = rx.reconcileExport()
	case cloudv1.ExportPhaseCompleted, cloudv1.ExportPhaseFailed:
		err = rx.cleanExport()
	case cloudv1.ExportPhaseCancel:
		err = rx.cleanExport()
	}
	if err != nil {
		if err == utils.ErrKeyspaceMetaServiceGroupMigrating {
			return ctrl.Result{RequeueAfter: requeueAfter}, nil
		}
		return utils.ReconcileErr(log, "Export", rx.export.Name, err)
	}

	// remove finalizer should be the last step (resource need to be cleaned).
	if rx.export.DeletionTimestamp != nil {
		// only allow to delete when export is in final status
		if exportPhase == cloudv1.ExportPhaseFailed {
			rx.Info("try remove finalizer")
			err = rx.tryRemoveFinalizer()
		}
		// only allow to delete when export is in final status and metering is done
		if (exportPhase == cloudv1.ExportPhaseCancel || exportPhase == cloudv1.ExportPhaseCompleted) &&
			rx.export.Annotations[types.AnnoExportConsoleMetric] != "" {
			rx.Info("try remove finalizer")
			err = rx.tryRemoveFinalizer()
		}
	}
	return utils.ReconcileErr(log, "Export", rx.export.Name, err)
}

// SetupWithManager sets up the controller with the Manager.
//
//nolint:mnd,gosec // ignore
func (r *Reconciler) SetupWithManager(exportCRMgr ctrlmanager.Manager, controllerName string) error {
	r.initRuntimeWithExportCluster(exportCRMgr)
	r.Recorder = exportCRMgr.GetEventRecorderFor("Export")
	r.HttpClient = &http.Client{
		Timeout: ExportAssignTimeout,
		Transport: &http.Transport{
			IdleConnTimeout:     60 * time.Second,
			TLSHandshakeTimeout: 5 * time.Second,
			TLSClientConfig:     &tls.Config{InsecureSkipVerify: true},
			ForceAttemptHTTP2:   true,
		},
	}

	c, err := controller.New(controllerName, exportCRMgr, controller.Options{
		MaxConcurrentReconciles: maxConcurrentExports(),
		Reconciler:              r,
	})
	if err != nil {
		return fmt.Errorf("can't new a controller: %w", err)
	}

	if err := exportCRMgr.GetFieldIndexer().IndexField(context.TODO(), &cloudv1.Export{},
		"spec.clusterName", func(obj client.Object) []string {
			export := obj.(*cloudv1.Export)
			clusterName := export.Spec.ClusterName
			if clusterName != "" {
				return []string{clusterName}
			}
			return nil
		}); err != nil {
		return fmt.Errorf("unable to add index for export clusterName: %w", err)
	}
	return r.watches(c, exportCRMgr)
}

func maxConcurrentExports() int {
	maxConcurrent := MaxConcurrentExport
	if value := os.Getenv(maxConcurrentExportsEnv); value != "" {
		if configured, err := strconv.Atoi(value); err == nil && configured > 0 {
			maxConcurrent = configured
		}
	}
	return maxConcurrent
}

// jobLabelSelectorPredicate creates a TypedPredicate for *batchv1.Job using a LabelSelector.
func jobLabelSelectorPredicate(selector metav1.LabelSelector) (predicate.TypedPredicate[*batchv1.Job], error) {
	labelSelector, err := metav1.LabelSelectorAsSelector(&selector)
	if err != nil {
		return nil, err
	}
	return predicate.NewTypedPredicateFuncs(func(job *batchv1.Job) bool {
		return labelSelector.Matches(k8slabels.Set(job.GetLabels()))
	}), nil
}

func (r *Reconciler) watches(c controller.Controller, exportCRMgr ctrlmanager.Manager) error {
	if err := c.Watch(source.Kind(exportCRMgr.GetCache(), &cloudv1.Export{},
		&handler.TypedEnqueueRequestForObject[*cloudv1.Export]{})); err != nil {
		return fmt.Errorf("can't watch export: %w", err)
	}

	pre, err := jobLabelSelectorPredicate(metav1.LabelSelector{
		MatchLabels: map[string]string{
			types.LabelJobType:               DumplingJobType,
			cloudv1.ProviderTypeLabel:        r.ProviderType,
			labels.LabelKeyKubernetesCluster: r.KubeClusterName,
		},
	})
	if err != nil {
		return fmt.Errorf("invalid label selector: %w", err)
	}

	if err := c.Watch(source.Kind(
		r.WorkloadCache,
		&batchv1.Job{},
		handler.TypedFuncs[*batchv1.Job, ctrl.Request]{UpdateFunc: r.handleUpdateDumplingJobEvent},
		pre,
	)); err != nil {
		return fmt.Errorf("can't watch dumpling job")
	}
	return nil
}

func (r *Reconciler) decodeExportJob(ctx context.Context) (*batchv1.Job, error) {
	exportJobTemplateCm := &corev1.ConfigMap{}
	if err := r.Client.Get(ctx,
		client.ObjectKey{
			Name:      r.ExportJobCM,
			Namespace: r.ExportJobNS,
		}, exportJobTemplateCm); err != nil {
		return nil, err
	}
	jobTemplate := exportJobTemplateCm.Data[r.ExportJobCM]
	if jobTemplate == "" {
		return nil, errors.New("exportJobTemplate is empty")
	}
	exportJob := &batchv1.Job{}
	if err := runtime.DecodeInto(r.Codecs.UniversalDecoder(), []byte(jobTemplate), exportJob); err != nil {
		return nil, err
	}
	return exportJob, nil
}

func (r *Reconciler) getExportJob(ctx context.Context) (*batchv1.Job, error) {
	// The steady-state path is a lock-free load. CompareAndSwap is only used
	// when concurrent reconciles race to publish the first decoded template.
	if exportJob := r.exportJob.Load(); exportJob != nil {
		return exportJob, nil
	}

	exportJob, err := r.decodeExportJob(ctx)
	if err != nil {
		// Do not cache transient ConfigMap read or decode failures. A later
		// reconcile can retry initialization.
		return nil, err
	}
	if r.exportJob.CompareAndSwap(nil, exportJob) {
		return exportJob, nil
	}
	return r.exportJob.Load(), nil
}

// TODO distinguish vip and super-vip
func (r *reconcileContext) applyCfgFromCm(ctx context.Context, exportJob *batchv1.Job) (*batchv1.Job, error) {
	jo := exportJob.DeepCopy()
	// get export cm
	exportCm := &corev1.ConfigMap{}
	if err := r.Client.Get(ctx,
		client.ObjectKey{
			Name:      ExportConfigCM,
			Namespace: r.ExportJobNS,
		}, exportCm); err != nil {
		return nil, err
	}

	dumplingContainerIndex := -1
	for i, c := range jo.Spec.Template.Spec.Containers {
		if c.Name == DumplingContainerName {
			dumplingContainerIndex = i
			break
		}
	}
	if dumplingContainerIndex == -1 {
		return nil, errors.New("dumpling container not found")
	}

	isExportToParuqet := r.export.Spec.Options.FileType == "PARQUET"

	exportType := r.export.Annotations[types.AnnoExportType]

	cpuRequestKey := cpuResourceRequestKey
	cpuLimitKey := cpuResourceLimitKey
	memoryRequestKey := memoryResourceRequestKey
	memoryLimitKey := memoryResourceLimitKey
	if isExportToParuqet {
		memoryRequestKey = memoryResourceRequestParquetKey
		memoryLimitKey = memoryResourceLimitParquetKey
	}
	dumplingConcurrencyKey := dumplingConcurrencyKey
	nodeSelectorKey := nodeSelector
	tolerationsKey := tolerations
	if exportType == string(ExportTypeVIP) {
		cpuRequestKey = vipCpuResourceRequestKey
		cpuLimitKey = vipCpuResourceLimitKey
		if isExportToParuqet {
			memoryRequestKey = vipMemoryResourceRequestParquetKey
			memoryLimitKey = vipMemoryResourceLimitParquetKey
		} else {
			memoryRequestKey = vipMemoryResourceRequestKey
			memoryLimitKey = vipMemoryResourceLimitKey
		}
		dumplingConcurrencyKey = vipDumplingConcurrencyKey
		nodeSelectorKey = vipNodeSelector
		tolerationsKey = vipTolerations
	} else if exportType == string(ExportTypeSuperVIP) {
		cpuRequestKey = superVipCpuResourceRequestKey
		cpuLimitKey = superVipCpuResourceLimitKey
		memoryRequestKey = superVipMemoryResourceRequestKey
		memoryLimitKey = superVipMemoryResourceLimitKey
		dumplingConcurrencyKey = superVipDumplingConcurrencyKey
		nodeSelectorKey = superVipNodeSelector
		tolerationsKey = superVipTolerations
	}

	// resource
	containers := jo.Spec.Template.Spec.Containers
	if val, ok := exportCm.Data[cpuRequestKey]; ok {
		containers[dumplingContainerIndex].Resources.Requests[corev1.ResourceCPU] = resource.MustParse(val)
	}
	if val, ok := exportCm.Data[cpuLimitKey]; ok {
		containers[dumplingContainerIndex].Resources.Limits[corev1.ResourceCPU] = resource.MustParse(val)
	}
	if val, ok := exportCm.Data[memoryRequestKey]; ok {
		containers[dumplingContainerIndex].Resources.Requests[corev1.ResourceMemory] = resource.MustParse(val)
	}
	if val, ok := exportCm.Data[memoryLimitKey]; ok {
		containers[dumplingContainerIndex].Resources.Limits[corev1.ResourceMemory] = resource.MustParse(val)
	}
	// command
	if val, ok := exportCm.Data[dumplingConcurrencyKey]; ok {
		command := fmt.Sprintf("-t=%s", val)
		containers[dumplingContainerIndex].Command = append(containers[dumplingContainerIndex].Command, command)
	}
	// nodeSelector and tolerations
	if val, ok := exportCm.Data[nodeSelectorKey]; ok {
		jo.Spec.Template.Spec.NodeSelector =
			utils.UpdateMap(jo.Spec.Template.Spec.NodeSelector, labels.ServerlessComponentNodeLabelKey, val)
	}
	if val, ok := exportCm.Data[tolerationsKey]; ok {
		jo.Spec.Template.Spec.Tolerations = append(jo.Spec.Template.Spec.Tolerations, corev1.Toleration{
			Key:      labels.ServerlessComponentNodeLabelKey,
			Operator: corev1.TolerationOpEqual,
			Effect:   corev1.TaintEffectNoSchedule,
			Value:    val,
		})
	}
	return jo, nil
}

func (r *Reconciler) handleUpdateDumplingJobEvent(
	ctx context.Context,
	e event.TypedUpdateEvent[*batchv1.Job],
	q workqueue.TypedRateLimitingInterface[ctrl.Request],
) {
	if e.ObjectNew == nil {
		r.Logger.Error(nil, "UpdateEvent received with no metadata", "event", e)
		return
	}
	// enqueue the export that owns the dumpling job
	exportJob := e.ObjectNew
	ownedBy, namespace, name, _ := labels.GetOwnerFromLabels(exportJob.GetLabels())
	if ownedBy != labels.GVKString(&ExportGVK) {
		return
	}
	q.Add(ctrl.Request{
		NamespacedName: apitypes.NamespacedName{
			Namespace: namespace,
			Name:      name,
		},
	})
}