package export

import (
	"context"
	"io"
	"net/http"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/stretchr/testify/require"
	"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"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime/serializer"
	"k8s.io/client-go/tools/record"
	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/client/fake"
	"sigs.k8s.io/controller-runtime/pkg/cluster"

	providerscheme "github.com/tidbcloud/aws-shared-provider/pkg/scheme"
)

type fakeCache struct {
	ctrlcache.Cache
}

type roundTripperFunc func(*http.Request) (*http.Response, error)

func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
	return f(req)
}

func TestGetExportJobConcurrentLazyLoad(t *testing.T) {
	t.Parallel()

	const template = `apiVersion: batch/v1
kind: Job
metadata:
  name: export-template
spec:
  template:
    spec:
      restartPolicy: Never
      containers:
      - name: dumpling
        image: dumpling
`
	cm := &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{Name: "export-job-template", Namespace: "export-job"},
		Data:       map[string]string{"export-job-template": template},
	}
	r := &Reconciler{
		Client:      fake.NewClientBuilder().WithScheme(providerscheme.Scheme).WithObjects(cm).Build(),
		ExportJobCM: "export-job-template",
		ExportJobNS: "export-job",
		Codecs:      serializer.NewCodecFactory(providerscheme.Scheme),
	}

	const workers = 100
	jobs := make(chan *batchv1.Job, workers)
	errs := make(chan error, workers)
	var wg sync.WaitGroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			job, err := r.getExportJob(context.Background())
			jobs <- job
			errs <- err
		}()
	}
	wg.Wait()
	close(jobs)
	close(errs)

	for err := range errs {
		require.NoError(t, err)
	}
	for job := range jobs {
		require.Same(t, r.exportJob.Load(), job)
	}
}

func TestCleanExportRequeuesWhileWaitingForMetrics(t *testing.T) {
	t.Parallel()

	rx := &reconcileContext{
		Reconciler: &Reconciler{},
		ctx:        context.Background(),
		export: &cloudv1.Export{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{
			types.AnnoExportFinishTime: time.Now().Format(time.RFC3339),
		}}},
	}

	err := rx.cleanExport()
	requeueErr := &utils.RequeueAfterError{}
	require.ErrorAs(t, err, &requeueErr)
	require.Positive(t, requeueErr.RequeueAfter)
	require.LessOrEqual(t, requeueErr.RequeueAfter, time.Minute)
}

type fakeCluster struct {
	cluster.Cluster
	client   client.Client
	reader   client.Reader
	cache    ctrlcache.Cache
	recorder record.EventRecorder
}

func (f *fakeCluster) GetClient() client.Client {
	return f.client
}

func (f *fakeCluster) GetAPIReader() client.Reader {
	return f.reader
}

func (f *fakeCluster) GetCache() ctrlcache.Cache {
	return f.cache
}

func (f *fakeCluster) GetEventRecorderFor(string) record.EventRecorder {
	return f.recorder
}

func TestInitRuntimeWithExportClusterPreservesPreconfiguredWorkloadDeps(t *testing.T) {
	t.Parallel()

	workloadClient := fake.NewClientBuilder().WithScheme(providerscheme.Scheme).Build()
	workloadCache := &fakeCache{}
	exportClient := fake.NewClientBuilder().WithScheme(providerscheme.Scheme).Build()
	exportRecorder := record.NewFakeRecorder(16)
	exportCluster := &fakeCluster{
		client:   exportClient,
		reader:   exportClient,
		cache:    &fakeCache{},
		recorder: exportRecorder,
	}

	r := &Reconciler{
		Client:        workloadClient,
		WorkloadCache: workloadCache,
	}

	r.initRuntimeWithExportCluster(exportCluster)

	require.Same(t, workloadClient, r.Client)
	require.Same(t, workloadCache, r.WorkloadCache)
	require.Same(t, exportClient, r.InfraClient)
	require.Same(t, exportClient, r.InfraAPIReader)
	require.Same(t, exportRecorder, r.InfraRecorder)
}

func TestNewReconcileContextReadsExportFromConfiguredExportClient(t *testing.T) {
	t.Parallel()

	exportObj := &cloudv1.Export{}
	exportObj.Name = "export-1"
	exportObj.Namespace = "ns-1"
	exportObj.Labels = map[string]string{
		labels.LabelKeyKubernetesCluster: "same-cluster",
	}
	exportObj.Annotations = map[string]string{
		types.AnnoExportMeta: `{"projectId":"project-1","tenantId":"tenant-1","clusterId":"cluster-1","keyspaceId":1,"keyspaceToken":"ks-1","tidbPoolSelector":"default"}`,
	}

	exportClient := fake.NewClientBuilder().
		WithScheme(providerscheme.Scheme).
		WithRuntimeObjects(exportObj).
		Build()

	r := &Reconciler{
		KubeClusterName: "same-cluster",
	}
	r.initRuntimeWithExportCluster(&fakeCluster{
		client:   exportClient,
		reader:   exportClient,
		cache:    &fakeCache{},
		recorder: record.NewFakeRecorder(16),
	})

	rx, err := r.newReconcileContext(context.Background(), ctrl.Request{
		NamespacedName: client.ObjectKey{Namespace: exportObj.Namespace, Name: exportObj.Name},
	})
	require.NoError(t, err)
	require.NotNil(t, rx)
	require.Equal(t, exportObj.Name, rx.export.Name)
}

func TestTryCleanDumplingJobRemovesPodFinalizerWhenDeleteJobSucceeds(t *testing.T) {
	t.Parallel()

	ctx := context.Background()
	exportObj := &cloudv1.Export{}
	exportObj.Name = "export-1"
	exportObj.Namespace = "ns-1"
	exportObj.Status.Phase = cloudv1.ExportPhaseFailed

	job := &batchv1.Job{
		ObjectMeta: metav1.ObjectMeta{
			Name:      exportObj.Name + "-dumpling",
			Namespace: "export-job",
		},
	}
	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:       "export-1-dumpling-pod",
			Namespace:  "export-job",
			Labels:     map[string]string{"job-name": job.Name},
			Finalizers: []string{exportFinalizer},
		},
		Status: corev1.PodStatus{
			Phase: corev1.PodSucceeded,
		},
	}

	workloadClient := fake.NewClientBuilder().
		WithScheme(providerscheme.Scheme).
		WithRuntimeObjects(job, pod).
		Build()
	rx := &reconcileContext{
		Reconciler: &Reconciler{
			Client:      workloadClient,
			ExportJobNS: "export-job",
		},
		ctx:    ctx,
		export: exportObj,
	}

	require.NoError(t, rx.tryCleanDumplingJob(ctx))

	cleanedPod := &corev1.Pod{}
	require.NoError(t, workloadClient.Get(ctx, client.ObjectKeyFromObject(pod), cleanedPod))
	require.NotContains(t, cleanedPod.Finalizers, exportFinalizer)
}

func TestDumplingFileTypeUsesLowercase(t *testing.T) {
	t.Parallel()

	require.Equal(t, "csv", dumplingFileType(cloudv1.FileTypeCSV))
	require.Equal(t, "sql", dumplingFileType(cloudv1.FileTypeSQL))
	require.Equal(t, "parquet", dumplingFileType(cloudv1.FileType("PARQUET")))
}

func TestBuildDumplingColumnFilter(t *testing.T) {
	t.Parallel()

	filter, err := buildDumplingColumnFilter(cloudv1.ColumnSelector{
		Matcher: []string{"app.orders", "app.customers"},
		Columns: []string{"*", "!customer_email"},
	})
	require.NoError(t, err)
	require.Equal(t, `{ matcher = ["app.orders","app.customers"], columns = ["*","!customer_email"] }`, filter)

	var parsed struct {
		Filter struct {
			Matcher []string `toml:"matcher"`
			Columns []string `toml:"columns"`
		} `toml:"filter"`
	}
	_, err = toml.Decode("filter = "+filter, &parsed)
	require.NoError(t, err)
	require.Equal(t, []string{"app.orders", "app.customers"}, parsed.Filter.Matcher)
	require.Equal(t, []string{"*", "!customer_email"}, parsed.Filter.Columns)

	_, err = buildDumplingColumnFilter(cloudv1.ColumnSelector{Columns: []string{"*"}})
	require.Error(t, err)
	_, err = buildDumplingColumnFilter(cloudv1.ColumnSelector{Matcher: []string{"app.orders"}})
	require.Error(t, err)
}

func TestMaxConcurrentExports(t *testing.T) {
	tests := []struct {
		name  string
		value string
		want  int
	}{
		{name: "default", want: MaxConcurrentExport},
		{name: "configured", value: "25", want: 25},
		{name: "zero uses default", value: "0", want: MaxConcurrentExport},
		{name: "negative uses default", value: "-1", want: MaxConcurrentExport},
		{name: "invalid uses default", value: "invalid", want: MaxConcurrentExport},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Setenv(maxConcurrentExportsEnv, tt.value)
			require.Equal(t, tt.want, maxConcurrentExports())
		})
	}
}

func TestEnsureTiDBPodDoesNotMutatePoolConfig(t *testing.T) {
	t.Parallel()

	ctx := context.Background()
	exportObj := &cloudv1.Export{
		ObjectMeta: metav1.ObjectMeta{
			Name:        "export-1",
			Namespace:   "ns-1",
			Annotations: map[string]string{},
		},
	}
	infraClient := fake.NewClientBuilder().WithScheme(providerscheme.Scheme).WithObjects(exportObj.DeepCopy()).Build()
	reconciledExport := &cloudv1.Export{}
	require.NoError(t, infraClient.Get(ctx, client.ObjectKeyFromObject(exportObj), reconciledExport))
	reconciledExport.Annotations = map[string]string{}

	poolConfig := &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{Name: "serverless-pool-config", Namespace: ExportTiDBPoolNS},
		Data: map[string]string{
			types.MinPoolReplicasKey: "0",
			types.MaxPoolReplicasKey: "3",
		},
	}
	workloadClient := fake.NewClientBuilder().WithScheme(providerscheme.Scheme).WithObjects(poolConfig).Build()
	httpClient := &http.Client{Transport: roundTripperFunc(func(*http.Request) (*http.Response, error) {
		return &http.Response{
			StatusCode: http.StatusOK,
			Status:     http.StatusText(http.StatusOK),
			Body:       io.NopCloser(strings.NewReader("pod-1,10.0.0.1")),
		}, nil
	})}

	rx := &reconcileContext{
		Reconciler: &Reconciler{
			Client:      workloadClient,
			InfraClient: infraClient,
			HttpClient:  httpClient,
		},
		ctx:           ctx,
		export:        reconciledExport,
		keyspaceToken: "keyspace-1",
	}

	require.NoError(t, rx.ensureTiDBPod())
	require.Equal(t, "pod-1,10.0.0.1", rx.export.Annotations[types.AnnoExportTiDBPod])

	pool := &corev1.ConfigMap{}
	require.NoError(t, workloadClient.Get(ctx, client.ObjectKeyFromObject(poolConfig), pool))
	require.Equal(t, "0", pool.Data[types.MinPoolReplicasKey])
	require.Equal(t, "3", pool.Data[types.MaxPoolReplicasKey])
}

// newAssignFailureContext builds a reconcile context whose assign call always fails, so the
// on-demand pod creation path can be exercised without a running manager server.
func newAssignFailureContext(t *testing.T, createdAt time.Time, exportCms ...client.Object) *reconcileContext {
	t.Helper()

	ctx := context.Background()
	exportObj := &cloudv1.Export{
		ObjectMeta: metav1.ObjectMeta{
			Name:              "export-1",
			Namespace:         "ns-1",
			CreationTimestamp: metav1.NewTime(createdAt),
			Annotations:       map[string]string{},
		},
	}
	httpClient := &http.Client{Transport: roundTripperFunc(func(*http.Request) (*http.Response, error) {
		return &http.Response{
			StatusCode: http.StatusInternalServerError,
			Status:     http.StatusText(http.StatusInternalServerError),
			Body:       io.NopCloser(strings.NewReader("not enough pod for assign")),
		}, nil
	})}

	infraClient := fake.NewClientBuilder().WithScheme(providerscheme.Scheme).
		WithStatusSubresource(&cloudv1.Export{}).
		WithObjects(exportObj.DeepCopy()).Build()
	// Read the export back so status updates carry a resource version the fake client accepts.
	reconciledExport := &cloudv1.Export{}
	require.NoError(t, infraClient.Get(ctx, client.ObjectKeyFromObject(exportObj), reconciledExport))
	reconciledExport.Annotations = map[string]string{}

	return &reconcileContext{
		Reconciler: &Reconciler{
			Client: fake.NewClientBuilder().WithScheme(providerscheme.Scheme).
				WithObjects(exportCms...).Build(),
			InfraClient:   infraClient,
			InfraRecorder: record.NewFakeRecorder(10),
			HttpClient:    httpClient,
			ExportJobNS:   "export-job",
		},
		ctx:           ctx,
		export:        reconciledExport,
		keyspaceToken: "keyspace-1",
	}
}

func TestEnsureTiDBPodRequeuesWhileWaitingForOnDemandPod(t *testing.T) {
	t.Parallel()

	rx := newAssignFailureContext(t, time.Now())

	err := rx.ensureTiDBPod()
	requeueErr := &utils.RequeueAfterError{}
	require.ErrorAs(t, err, &requeueErr)
	require.Equal(t, assignTiDBRequeueInterval, requeueErr.RequeueAfter)
	require.Empty(t, rx.export.Annotations[types.AnnoExportTiDBPod])
}

func TestEnsureTiDBPodFailsAfterMaxWait(t *testing.T) {
	t.Parallel()

	rx := newAssignFailureContext(t, time.Now().Add(-DefaultAssignTiDBMaxWait-time.Minute))

	err := rx.ensureTiDBPod()
	unRetryableErr := &UnRetryableError{}
	require.ErrorAs(t, err, &unRetryableErr)
	require.Contains(t, unRetryableErr.Error(), ExportTiDBPoolNS)
	require.Contains(t, unRetryableErr.Error(), DefaultAssignTiDBMaxWait.String())

	requeueErr := &utils.RequeueAfterError{}
	require.NotErrorAs(t, err, &requeueErr)
}

// Exhausting the wait budget must reach a terminal phase, otherwise the export keeps
// retrying forever and never releases its etcd token and resource group.
func TestStartExportMarksExportFailedAfterMaxWait(t *testing.T) {
	t.Parallel()

	rx := newAssignFailureContext(t, time.Now().Add(-DefaultAssignTiDBMaxWait-time.Minute))
	rx.export.Annotations = map[string]string{
		types.AnnoExportType:          string(ExportTypeVIP),
		types.AnnoExportEtcdKey:       "etcd-key",
		types.AnnoExportResourceGroup: "resource-group",
	}

	require.NoError(t, rx.startExport())
	require.Equal(t, cloudv1.ExportPhaseFailed, rx.export.Status.Phase)
	require.Contains(t, rx.export.Status.Message, ExportVipTiDBPoolNS)

	persisted := &cloudv1.Export{}
	require.NoError(t, rx.InfraClient.Get(rx.ctx, client.ObjectKeyFromObject(rx.export), persisted))
	require.Equal(t, cloudv1.ExportPhaseFailed, persisted.Status.Phase)
}

// startExport must not wrap the requeue error, otherwise resolveReconcileErr degrades it
// into a plain reconcile failure.
func TestStartExportPropagatesAssignRequeue(t *testing.T) {
	t.Parallel()

	rx := newAssignFailureContext(t, time.Now())
	rx.export.Annotations = map[string]string{
		types.AnnoExportType:          string(ExportTypeVIP),
		types.AnnoExportEtcdKey:       "etcd-key",
		types.AnnoExportResourceGroup: "resource-group",
	}

	err := rx.startExport()
	requeueErr := &utils.RequeueAfterError{}
	require.ErrorAs(t, err, &requeueErr)

	result, reconcileErr := utils.ReconcileErr(rx.Logger, "Export", rx.export.Name, err)
	require.NoError(t, reconcileErr)
	require.Equal(t, assignTiDBRequeueInterval, result.RequeueAfter)
}

func TestAssignTiDBMaxWait(t *testing.T) {
	tests := []struct {
		name  string
		value string
		want  time.Duration
	}{
		{name: "no config map", want: DefaultAssignTiDBMaxWait},
		{name: "empty uses default", value: "", want: DefaultAssignTiDBMaxWait},
		{name: "configured", value: "2m", want: 2 * time.Minute},
		{name: "zero uses default", value: "0s", want: DefaultAssignTiDBMaxWait},
		{name: "negative uses default", value: "-1m", want: DefaultAssignTiDBMaxWait},
		{name: "invalid uses default", value: "invalid", want: DefaultAssignTiDBMaxWait},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			var exportCms []client.Object
			if tt.name != "no config map" {
				exportCms = append(exportCms, &corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: ExportConfigCM, Namespace: "export-job"},
					Data:       map[string]string{assignTiDBMaxWaitKey: tt.value},
				})
			}

			rx := newAssignFailureContext(t, time.Now(), exportCms...)
			require.Equal(t, tt.want, rx.assignTiDBMaxWait())
		})
	}
}