package zonalshift

import (
	"context"
	"fmt"
	"testing"

	"github.com/onsi/gomega"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	cloudaws "github.com/tidbcloud/infra-provider/pkg/cloud/aws"
	"github.com/tidbcloud/infra-provider/pkg/dataplane"
)

const testNLBARN = "arn:aws:elasticloadbalancing:us-east-1:123:loadbalancer/net/nlb/abc"
const testTGARN = "arn:aws:elasticloadbalancing:us-east-1:123:targetgroup/tg/xyz"

type fakeAWSAPI struct {
	resolveAZIDFn                   func(ctx context.Context, role *cloudaws.AssumedRole, region, azName string) (string, error)
	describeNLBFn                   func(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string) (privateNLB, error)
	ensureTGConnectionTerminationFn func(ctx context.Context, role *cloudaws.AssumedRole, region string, tgARNs []string, enabled bool) error
	ensureNLBZonalShiftFn           func(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string, enabled bool) error
	startZonalShiftFn               func(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN, azID, expiresIn, comment string) (string, error)
	updateZonalShiftFn              func(ctx context.Context, role *cloudaws.AssumedRole, region, shiftID, expiresIn string) error
	listActiveZonalShiftsFn         func(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string) ([]activeZonalShift, error)
	cancelZonalShiftFn              func(ctx context.Context, role *cloudaws.AssumedRole, region, shiftID string) error

	tgEnabledCalls  []bool
	nlbEnabledCalls []bool
	cancelledIDs    []string
	startCalls      int
	describedARNs   []string
}

func (f *fakeAWSAPI) ResolveAZID(ctx context.Context, role *cloudaws.AssumedRole, region, azName string) (string, error) {
	if f.resolveAZIDFn != nil {
		return f.resolveAZIDFn(ctx, role, region, azName)
	}
	return "use1-az1", nil
}

func (f *fakeAWSAPI) DescribeNLB(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string) (privateNLB, error) {
	f.describedARNs = append(f.describedARNs, nlbARN)
	if f.describeNLBFn != nil {
		return f.describeNLBFn(ctx, role, region, nlbARN)
	}
	return privateNLB{ARN: nlbARN, TGARNs: []string{testTGARN}}, nil
}

func (f *fakeAWSAPI) EnsureTGConnectionTermination(ctx context.Context, role *cloudaws.AssumedRole, region string, tgARNs []string, enabled bool) error {
	f.tgEnabledCalls = append(f.tgEnabledCalls, enabled)
	if f.ensureTGConnectionTerminationFn != nil {
		return f.ensureTGConnectionTerminationFn(ctx, role, region, tgARNs, enabled)
	}
	return nil
}

func (f *fakeAWSAPI) EnsureNLBZonalShift(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string, enabled bool) error {
	f.nlbEnabledCalls = append(f.nlbEnabledCalls, enabled)
	if f.ensureNLBZonalShiftFn != nil {
		return f.ensureNLBZonalShiftFn(ctx, role, region, nlbARN, enabled)
	}
	return nil
}

func (f *fakeAWSAPI) StartZonalShift(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN, azID, expiresIn, comment string) (string, error) {
	f.startCalls++
	if f.startZonalShiftFn != nil {
		return f.startZonalShiftFn(ctx, role, region, nlbARN, azID, expiresIn, comment)
	}
	return "shift-123", nil
}

func (f *fakeAWSAPI) UpdateZonalShift(ctx context.Context, role *cloudaws.AssumedRole, region, shiftID, expiresIn string) error {
	if f.updateZonalShiftFn != nil {
		return f.updateZonalShiftFn(ctx, role, region, shiftID, expiresIn)
	}
	return nil
}

func (f *fakeAWSAPI) ListActiveZonalShifts(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string) ([]activeZonalShift, error) {
	if f.listActiveZonalShiftsFn != nil {
		return f.listActiveZonalShiftsFn(ctx, role, region, nlbARN)
	}
	return nil, nil
}

func (f *fakeAWSAPI) CancelZonalShift(ctx context.Context, role *cloudaws.AssumedRole, region, shiftID string) error {
	f.cancelledIDs = append(f.cancelledIDs, shiftID)
	if f.cancelZonalShiftFn != nil {
		return f.cancelZonalShiftFn(ctx, role, region, shiftID)
	}
	return nil
}

type stubEnvProvider struct {
	dataplane.NopEnvProvider
	region  string
	account string
}

func (s *stubEnvProvider) GetRegion() string  { return s.region }
func (s *stubEnvProvider) GetAccount() string { return s.account }

func newTestCluster() *cloudv1.Cluster {
	return &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "c1",
			Namespace: "ns1",
			Labels: map[string]string{
				cloudv1.ProviderTypeLabel: cloudv1.ProviderTypeEKS,
				cloudv1.ClusterTagKey:     "c1",
				cloudv1.RegionLabel:       "us-east-1",
			},
			Annotations: map[string]string{
				cloudv1.DefaultLoadBalancerIDAnnotationKey: "default-gid",
			},
		},
		Spec: cloudv1.ClusterSpec{
			Kubernetes: cloudv1.Kubernetes{
				Provider: cloudv1.ProviderTypeEKS,
			},
		},
		Status: cloudv1.ClusterStatus{
			SqlEndpoints: cloudv1.SqlEndpointsStatus{
				PrivateEndpoint: &cloudv1.SqlEndpointStatus{
					Host:       "private.example.com",
					Port:       4000,
					ResourceID: testNLBARN,
				},
			},
		},
	}
}

func newEnvMgr(scheme *runtime.Scheme) dataplane.EnvProviderManager {
	return dataplane.NewFakeEnvProviderManager(scheme, func(ctx context.Context, scheme *runtime.Scheme, c *cloudv1.Cluster) (dataplane.EnvProvider, error) {
		return &stubEnvProvider{region: "us-east-1", account: "123456789012"}, nil
	})
}

func TestStartZonalShiftHandle(t *testing.T) {
	g := gomega.NewWithT(t)
	scheme := runtime.NewScheme()
	g.Expect(cloudv1.AddToScheme(scheme)).To(gomega.Succeed())

	cluster := newTestCluster()
	wf := &cloudv1.Workflow{
		ObjectMeta: metav1.ObjectMeta{Name: "wf1", Namespace: "ns1"},
		Spec: cloudv1.WorkflowSpec{
			Steps: cloudv1.WorkflowSteps{{
				Name: "Start zonal shift",
				Inputs: cloudv1.InputsWrapper{Object: map[string]interface{}{
					"namespace":        "ns1",
					"clusterName":      "c1",
					"availabilityZone": "us-east-1a",
					"expiresIn":        "10m",
				}},
			}},
		},
		Status: cloudv1.WorkflowStatus{
			Phase: cloudv1.WorkflowPhaseRunning,
			Steps: []*cloudv1.WorkflowStepStatus{{
				Name:    "Start zonal shift",
				Phase:   cloudv1.WorkflowStepPhaseRunning,
				Outputs: map[string]string{},
			}},
		},
	}
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster, wf).WithStatusSubresource(&cloudv1.Workflow{}).Build()
	api := &fakeAWSAPI{}
	action := NewStartZonalShift(cli, api, newEnvMgr(scheme), "DBaaS-Role", "ext")

	g.Expect(action.Handle(context.Background(), wf, wf.Spec.Steps[0])).To(gomega.Succeed())
	g.Expect(api.tgEnabledCalls).To(gomega.Equal([]bool{false}))
	g.Expect(api.nlbEnabledCalls).To(gomega.Equal([]bool{true}))
	g.Expect(api.startCalls).To(gomega.Equal(1))
	g.Expect(wf.Status.Steps[0].Outputs[zonalShiftIDsOutputKey]).To(gomega.Equal("shift-123"))
}

func TestStartZonalShiftRetriesWhenNotManaged(t *testing.T) {
	g := gomega.NewWithT(t)
	scheme := runtime.NewScheme()
	g.Expect(cloudv1.AddToScheme(scheme)).To(gomega.Succeed())

	cluster := newTestCluster()
	wf := &cloudv1.Workflow{
		ObjectMeta: metav1.ObjectMeta{Name: "wf1", Namespace: "ns1"},
		Spec: cloudv1.WorkflowSpec{
			Steps: cloudv1.WorkflowSteps{{
				Name: "Start zonal shift",
				Inputs: cloudv1.InputsWrapper{Object: map[string]interface{}{
					"namespace": "ns1", "clusterName": "c1", "availabilityZone": "us-east-1a", "expiresIn": "10m",
				}},
			}},
		},
		Status: cloudv1.WorkflowStatus{
			Phase: cloudv1.WorkflowPhaseRunning,
			Steps: []*cloudv1.WorkflowStepStatus{{Name: "Start zonal shift", Phase: cloudv1.WorkflowStepPhaseRunning, Outputs: map[string]string{}}},
		},
	}
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster, wf).WithStatusSubresource(&cloudv1.Workflow{}).Build()
	api := &fakeAWSAPI{
		startZonalShiftFn: func(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN, azID, expiresIn, comment string) (string, error) {
			return "", fmt.Errorf("resource arn:... is not a managed resource")
		},
	}
	action := NewStartZonalShift(cli, api, newEnvMgr(scheme), "DBaaS-Role", "")
	err := action.Handle(context.Background(), wf, wf.Spec.Steps[0])
	g.Expect(err).To(gomega.HaveOccurred())
	g.Expect(err.Error()).To(gomega.ContainSubstring("will retry"))
	g.Expect(err.Error()).NotTo(gomega.ContainSubstring("non-retriable"))
}

func TestStartZonalShiftIdempotent(t *testing.T) {
	g := gomega.NewWithT(t)
	scheme := runtime.NewScheme()
	g.Expect(cloudv1.AddToScheme(scheme)).To(gomega.Succeed())

	cluster := newTestCluster()
	wf := &cloudv1.Workflow{
		ObjectMeta: metav1.ObjectMeta{Name: "wf1", Namespace: "ns1"},
		Spec: cloudv1.WorkflowSpec{
			Steps: cloudv1.WorkflowSteps{{
				Name: "Start zonal shift",
				Inputs: cloudv1.InputsWrapper{Object: map[string]interface{}{
					"namespace": "ns1", "clusterName": "c1", "availabilityZone": "us-east-1a", "expiresIn": "10m",
				}},
			}},
		},
		Status: cloudv1.WorkflowStatus{
			Phase: cloudv1.WorkflowPhaseRunning,
			Steps: []*cloudv1.WorkflowStepStatus{{Name: "Start zonal shift", Phase: cloudv1.WorkflowStepPhaseRunning, Outputs: map[string]string{}}},
		},
	}
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster, wf).WithStatusSubresource(&cloudv1.Workflow{}).Build()
	api := &fakeAWSAPI{
		listActiveZonalShiftsFn: func(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string) ([]activeZonalShift, error) {
			return []activeZonalShift{{ID: "existing-shift", AwayFrom: "use1-az1"}}, nil
		},
		startZonalShiftFn: func(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN, azID, expiresIn, comment string) (string, error) {
			t.Fatal("StartZonalShift should not be called when an active shift exists")
			return "", nil
		},
	}
	action := NewStartZonalShift(cli, api, newEnvMgr(scheme), "DBaaS-Role", "")
	g.Expect(action.Handle(context.Background(), wf, wf.Spec.Steps[0])).To(gomega.Succeed())
	g.Expect(api.startCalls).To(gomega.Equal(0))
	g.Expect(wf.Status.Steps[0].Outputs[zonalShiftIDsOutputKey]).To(gomega.Equal("existing-shift"))
}

func TestCancelZonalShiftIdempotent(t *testing.T) {
	g := gomega.NewWithT(t)
	scheme := runtime.NewScheme()
	g.Expect(cloudv1.AddToScheme(scheme)).To(gomega.Succeed())

	cluster := newTestCluster()
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster).Build()
	api := &fakeAWSAPI{}
	action := NewCancelZonalShift(cli, api, newEnvMgr(scheme), "DBaaS-Role", "")
	wf := &cloudv1.Workflow{
		Spec: cloudv1.WorkflowSpec{Steps: cloudv1.WorkflowSteps{{
			Name: "Cancel",
			Inputs: cloudv1.InputsWrapper{Object: map[string]interface{}{
				"namespace": "ns1", "clusterName": "c1",
			}},
		}}},
	}
	g.Expect(action.Handle(context.Background(), wf, wf.Spec.Steps[0])).To(gomega.Succeed())
	g.Expect(api.cancelledIDs).To(gomega.BeEmpty())
}

func TestCleanupZonalShiftOrder(t *testing.T) {
	g := gomega.NewWithT(t)
	scheme := runtime.NewScheme()
	g.Expect(cloudv1.AddToScheme(scheme)).To(gomega.Succeed())

	cluster := newTestCluster()
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster).Build()
	api := &fakeAWSAPI{
		listActiveZonalShiftsFn: func(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string) ([]activeZonalShift, error) {
			return []activeZonalShift{{ID: "s1", AwayFrom: "use1-az1"}}, nil
		},
	}
	action := NewCleanupZonalShift(cli, api, newEnvMgr(scheme), "DBaaS-Role", "")
	wf := &cloudv1.Workflow{
		Spec: cloudv1.WorkflowSpec{Steps: cloudv1.WorkflowSteps{{
			Name: "Cleanup",
			Inputs: cloudv1.InputsWrapper{Object: map[string]interface{}{
				"namespace": "ns1", "clusterName": "c1",
			}},
		}}},
	}
	g.Expect(action.Handle(context.Background(), wf, wf.Spec.Steps[0])).To(gomega.Succeed())
	g.Expect(api.cancelledIDs).To(gomega.Equal([]string{"s1"}))
	g.Expect(api.nlbEnabledCalls).To(gomega.Equal([]bool{false}))
	g.Expect(api.tgEnabledCalls).To(gomega.Equal([]bool{true}))
}

func TestUpdateZonalShiftHandle(t *testing.T) {
	g := gomega.NewWithT(t)
	scheme := runtime.NewScheme()
	g.Expect(cloudv1.AddToScheme(scheme)).To(gomega.Succeed())

	cluster := newTestCluster()
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster).Build()
	var gotIDs []string
	var gotExpires string
	api := &fakeAWSAPI{
		listActiveZonalShiftsFn: func(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string) ([]activeZonalShift, error) {
			return []activeZonalShift{
				{ID: "shift-abc", AwayFrom: "use1-az1"},
				{ID: "shift-def", AwayFrom: "use1-az1"},
			}, nil
		},
		updateZonalShiftFn: func(ctx context.Context, role *cloudaws.AssumedRole, region, shiftID, expiresIn string) error {
			gotIDs = append(gotIDs, shiftID)
			gotExpires = expiresIn
			return nil
		},
	}
	action := NewUpdateZonalShift(cli, api, newEnvMgr(scheme), "DBaaS-Role", "")
	wf := &cloudv1.Workflow{
		Spec: cloudv1.WorkflowSpec{Steps: cloudv1.WorkflowSteps{{
			Name: "Update",
			Inputs: cloudv1.InputsWrapper{Object: map[string]interface{}{
				"namespace": "ns1", "clusterName": "c1", "expiresIn": "30m",
			}},
		}}},
	}
	g.Expect(action.Handle(context.Background(), wf, wf.Spec.Steps[0])).To(gomega.Succeed())
	g.Expect(gotIDs).To(gomega.Equal([]string{"shift-abc", "shift-def"}))
	g.Expect(gotExpires).To(gomega.Equal("30m"))
}

func TestStartZonalShiftMultipleNLBs(t *testing.T) {
	g := gomega.NewWithT(t)
	scheme := runtime.NewScheme()
	g.Expect(cloudv1.AddToScheme(scheme)).To(gomega.Succeed())

	extraARN := "arn:aws:elasticloadbalancing:us-east-1:123:loadbalancer/net/extra/def"
	cluster := newTestCluster()
	cluster.Spec.ExtraGroups = []cloudv1.ExtraGroupSpec{{Name: "analytics", TiDB: cloudv1.TiDB{Replicas: 1}}}
	cluster.Status.ExtraEndpoints = []cloudv1.ExtraGroupStatus{{
		Name: "analytics",
		SqlEndpointsStatus: cloudv1.SqlEndpointsStatus{
			PrivateEndpoint: &cloudv1.SqlEndpointStatus{ResourceID: extraARN},
		},
	}}
	wf := &cloudv1.Workflow{
		ObjectMeta: metav1.ObjectMeta{Name: "wf1", Namespace: "ns1"},
		Spec: cloudv1.WorkflowSpec{
			Steps: cloudv1.WorkflowSteps{{
				Name: "Start zonal shift",
				Inputs: cloudv1.InputsWrapper{Object: map[string]interface{}{
					"namespace": "ns1", "clusterName": "c1", "availabilityZone": "us-east-1a", "expiresIn": "10m",
				}},
			}},
		},
		Status: cloudv1.WorkflowStatus{
			Phase: cloudv1.WorkflowPhaseRunning,
			Steps: []*cloudv1.WorkflowStepStatus{{Name: "Start zonal shift", Phase: cloudv1.WorkflowStepPhaseRunning, Outputs: map[string]string{}}},
		},
	}
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster, wf).WithStatusSubresource(&cloudv1.Workflow{}).Build()
	api := &fakeAWSAPI{
		describeNLBFn: func(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string) (privateNLB, error) {
			return privateNLB{ARN: nlbARN, TGARNs: []string{testTGARN}}, nil
		},
		startZonalShiftFn: func(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN, azID, expiresIn, comment string) (string, error) {
			return "shift-" + nlbARN[len(nlbARN)-3:], nil
		},
	}
	action := NewStartZonalShift(cli, api, newEnvMgr(scheme), "DBaaS-Role", "")
	g.Expect(action.Handle(context.Background(), wf, wf.Spec.Steps[0])).To(gomega.Succeed())
	g.Expect(api.describedARNs).To(gomega.Equal([]string{testNLBARN, extraARN}))
	g.Expect(api.nlbEnabledCalls).To(gomega.HaveLen(2))
	g.Expect(api.startCalls).To(gomega.Equal(2))
	g.Expect(wf.Status.Steps[0].Outputs[zonalShiftIDsOutputKey]).To(gomega.Equal("shift-abc,shift-def"))
}

func TestEnsureZonalShiftRejectsMismatchedAwayFrom(t *testing.T) {
	g := gomega.NewWithT(t)
	api := &fakeAWSAPI{
		listActiveZonalShiftsFn: func(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string) ([]activeZonalShift, error) {
			return []activeZonalShift{{ID: "shift-old", AwayFrom: "use1-az1"}}, nil
		},
		startZonalShiftFn: func(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN, azID, expiresIn, comment string) (string, error) {
			t.Fatal("StartZonalShift should not be called when an active shift exists for another AZ")
			return "", nil
		},
	}
	_, err := ensureZonalShift(context.Background(), api, nil, "us-east-1", testNLBARN, "use1-az2", "10m", "test")
	g.Expect(err).To(gomega.HaveOccurred())
	g.Expect(err.Error()).To(gomega.ContainSubstring("away from use1-az1"))
	g.Expect(err.Error()).To(gomega.ContainSubstring("requested use1-az2"))
}