package eks

import (
	"context"
	"fmt"
	"testing"
	"time"

	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	"github.com/crossplane/crossplane-runtime/pkg/meta"
	"github.com/go-logr/logr"
	"github.com/stretchr/testify/assert"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/upbound/provider-aws/apis/ec2/v1beta1"
	"go.uber.org/mock/gomock"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/yaml"

	"github.com/tidbcloud/infra-provider/regional-server/api/v1alpha1"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/config"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers/internet/task"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/deploy"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/k8s"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/namer"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/repository"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/util"
	"github.com/tidbcloud/infra-provider/regional-server/test/kubemock"
	"github.com/tidbcloud/infra-provider/regional-server/test/testdata"
)

func fakeTaskContext(t *testing.T) *task.Context {
	s := runtime.NewScheme()
	v1alpha1.AddToScheme(s)
	v1beta1.SchemeBuilder.AddToScheme(s)
	fakeVPC := &v1alpha1.VPC{}
	bs, err := testdata.LoadDefaultVPCYAML()
	assert.Nil(t, err)
	assert.Nil(t, yaml.Unmarshal(bs, fakeVPC))
	fakeVPC.Spec.Zones = []string{"zone1", "zone2", "zone3"}
	fakeInternet := &v1alpha1.Internet{}
	bs, err = testdata.LoadEKSInternetYAML()
	assert.Nil(t, err)
	assert.Nil(t, yaml.Unmarshal(bs, fakeInternet))
	cli := fake.NewClientBuilder().WithScheme(s).WithObjects(fakeVPC, fakeInternet).
		WithStatusSubresource(fakeVPC, fakeInternet, &v1beta1.InternetGateway{}, &v1beta1.EIP{}, &v1beta1.NATGateway{}).Build()
	c := &task.Context{
		Local:  k8s.NewMutClientWithOwner(cli, fakeVPC.DeepCopy(), namer.NewNamer()).Client(),
		Logger: logr.Discard(),
		VPC:    fakeVPC.DeepCopy(),
		VPCR:   repository.NewVPCReader(cli),
		Config: config.RootConfig{
			ControllerConfig: config.ControllerConfig{
				RegionConfig: config.RegionConfig{
					Name: "us-west-2",
					Zones: []string{
						"zone1",
						"zone2",
						"zone3",
					},
				},
			},
		},
		Internet:   fakeInternet.DeepCopy(),
		InternetRW: repository.NewInternetReaderWriter(cli),
	}
	c, err = c.Default()
	assert.Nil(t, err)
	return c
}

func fakeNATTask(t *testing.T) *NAT {
	return NewNAT(fakeTaskContext(t))
}

func TestNAT_Create(t *testing.T) {
	t.Run("create nat success", func(t *testing.T) {
		natTask := fakeNATTask(t)
		zones := natTask.GetNodeZones()
		natTask.VPC.Status.PublicSubnets = []v1alpha1.Subnet{}
		natTask.Internet.Status.IPAddresses = []v1alpha1.IPAddress{}
		name := natTask.InternetName()

		// subnet not found
		result := natTask.Create(context.Background())
		res, err, retriableErr := deploy.MinRequeue([]deploy.Result{result})
		assert.Nil(t, err)
		assert.ErrorContains(t, retriableErr, "no public subnet found")
		assert.Equal(t, time.Second*30, res.RequeueAfter)

		for _, zone := range zones {
			natTask.VPC.Status.PublicSubnets = append(natTask.VPC.Status.PublicSubnets, v1alpha1.Subnet{
				Name:               namer.PublicSubnetName(name, zone),
				AvailabilityZone:   zone,
				AvailabilityZoneID: fmt.Sprintf("%s-id", zone),
			})
			assert.Nil(t, natTask.Local.Create(context.Background(), &v1beta1.Subnet{
				ObjectMeta: metav1.ObjectMeta{
					Name: namer.PublicSubnetName(name, zone),
					Annotations: map[string]string{
						meta.AnnotationKeyExternalName: fmt.Sprintf("id-%s", zone),
					},
				},
			}))
			natTask.Internet.Status.IPAddresses = append(natTask.Internet.Status.IPAddresses, v1alpha1.IPAddress{
				Name:         namer.PublicSubnetName(name, zone),
				AllocationID: fmt.Sprintf("allo-%s", zone),
			})
		}

		assert.Nil(t, natTask.Local.Status().Update(context.Background(), natTask.VPC))
		internetIPAddress := natTask.Internet.Status.IPAddresses
		natTask.Internet.Status.IPAddresses = []v1alpha1.IPAddress{}
		result = natTask.Create(context.Background())
		res, err, retriableErr = deploy.MinRequeue([]deploy.Result{result})
		assert.Nil(t, err)
		assert.ErrorContains(t, retriableErr, "no ip address found for subnet")
		assert.Equal(t, time.Second*30, res.RequeueAfter)

		natTask.Internet.Status.IPAddresses = internetIPAddress
		result = natTask.Create(context.Background())
		assert.True(t, result.IsDoneState())
		assert.Equal(t, 3, len(natTask.nats))
	})
}

func TestNAT_createNAT(t *testing.T) {
	t.Run("create nat success", func(t *testing.T) {
		natTask := fakeNATTask(t)
		tags := cloudv1.TagsFromLabels(natTask.Internet.Labels)
		nat, err := natTask.createNAT(context.Background(), tags, "test", "zone", &v1beta1.Subnet{
			ObjectMeta: metav1.ObjectMeta{
				Name: "test",
				Annotations: map[string]string{
					meta.AnnotationKeyExternalName: "test-id",
				},
			},
		}, &v1alpha1.IPAddress{
			Name:         "test-ip",
			AllocationID: "test-allocation-id",
		})
		assert.Nil(t, err)
		assert.Equal(t, namer.PublicSubnetName("test", "zone"), nat.Name)
		assert.Equal(t, v1beta1.NATGatewaySpec{
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{Name: natTask.ProviderConfigName()},
				ManagementPolicies:      controllers.GetResourceManagementPolicies(),
			},
			ForProvider: v1beta1.NATGatewayParameters_2{
				Region:       pointer.String(natTask.Internet.Spec.Region),
				Tags:         util.ToMapPointer(tags),
				AllocationID: pointer.String("test-allocation-id"),
			},
			InitProvider: v1beta1.NATGatewayInitParameters_2{
				SubnetID:         pointer.String("test-id"),
				ConnectivityType: pointer.String("public"),
			},
		}, nat.Spec)
	})
	t.Run("create nat failed", func(t *testing.T) {
		natTask := fakeNATTask(t)
		s := natTask.Local.Scheme()
		ctl := gomock.NewController(t)
		mockcli := kubemock.NewMockClient(ctl)
		natTask.Local = mockcli
		mockcli.EXPECT().Scheme().Return(s).Times(1)
		mockcli.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(apierrors.NewNotFound(schema.GroupResource{}, "not found")).Times(1)
		mockcli.EXPECT().Create(gomock.Any(), gomock.Any()).Return(assert.AnError).Times(1)
		_, err := natTask.createNAT(context.Background(), nil, "test", "zone", &v1beta1.Subnet{}, &v1alpha1.IPAddress{})
		assert.ErrorIs(t, err, assert.AnError)
	})
}

func TestNAT_PostCreate(t *testing.T) {
	natTask := fakeNATTask(t)
	natTask.Config.ResourceNamePrefix = "tidbcloud"
	zones := natTask.GetNodeZones()
	natTask.VPC.Status.PublicSubnets = []v1alpha1.Subnet{}
	natTask.Internet.Status.IPAddresses = []v1alpha1.IPAddress{}
	name := natTask.InternetName()
	for _, zone := range zones {
		natTask.VPC.Status.PublicSubnets = append(natTask.VPC.Status.PublicSubnets, v1alpha1.Subnet{
			Name:               namer.PublicSubnetName(name, zone),
			AvailabilityZone:   zone,
			AvailabilityZoneID: fmt.Sprintf("%s-id", zone),
		})
		assert.Nil(t, natTask.Local.Create(context.Background(), &v1beta1.Subnet{
			ObjectMeta: metav1.ObjectMeta{
				Name: namer.PublicSubnetName(name, zone),
				Annotations: map[string]string{
					meta.AnnotationKeyExternalName: fmt.Sprintf("id-%s", zone),
				},
			},
		}))
		natTask.Internet.Status.IPAddresses = append(natTask.Internet.Status.IPAddresses, v1alpha1.IPAddress{
			Name:         namer.PublicSubnetName(name, zone),
			AllocationID: fmt.Sprintf("allo-%s", zone),
		})
	}
	assert.Nil(t, natTask.Local.Status().Update(context.Background(), natTask.VPC))
	result := natTask.Create(context.Background())
	assert.True(t, result.IsDoneState())
	assert.Equal(t, 3, len(natTask.nats))

	t.Run("nat not avtive", func(t *testing.T) {
		result = natTask.PostCreate(context.Background())
		res, err, retriableErr := deploy.MinRequeue([]deploy.Result{result})
		assert.Nil(t, err)
		assert.ErrorContains(t, retriableErr, "not available")
		assert.Equal(t, time.Second*30, res.RequeueAfter)
	})
	t.Run("nat avtive", func(t *testing.T) {
		nats := []v1alpha1.NAT{}
		for _, nat := range natTask.nats {
			natID := fmt.Sprintf("nat-%s", nat.Name)
			natRes := &v1beta1.NATGateway{}
			assert.Nil(t, natTask.Local.Get(context.Background(), types.NamespacedName{Name: nat.Name}, natRes))
			natRes.Status.SetConditions(xpv1.Available())
			natRes.Status.AtProvider.PublicIP = pointer.String("34.34.33.33")
			natRes.SetAnnotations(map[string]string{
				meta.AnnotationKeyExternalName: natID,
			})
			targetStatus := natRes.Status.DeepCopy()
			assert.Nil(t, natTask.Local.Update(context.Background(), natRes))
			natRes.Status = *targetStatus
			assert.Nil(t, natTask.Local.Status().Update(context.Background(), natRes))

			nats = append(nats, v1alpha1.NAT{
				Name:         nat.Name,
				NATGatewayID: natID,
			})
		}

		result = natTask.PostCreate(context.Background())
		assert.True(t, result.IsDoneState())
		assert.Equal(t, nats, natTask.Internet.Status.NATs)
	})
}

func Test_findPublicSubnetByZone(t *testing.T) {
	natTask := fakeNATTask(t)
	t.Run("subnet not found", func(t *testing.T) {
		sub, err := findPublicSubnetByZone(context.Background(), natTask.Local, natTask.VPC, "zone")
		assert.Nil(t, sub)
		assert.ErrorContains(t, err, "no public subnet found")
	})
	t.Run("subnet found", func(t *testing.T) {
		zones := natTask.GetNodeZones()
		name := natTask.InternetName()
		for _, zone := range zones {
			natTask.VPC.Status.PublicSubnets = append(natTask.VPC.Status.PublicSubnets, v1alpha1.Subnet{
				Name:               namer.PublicSubnetName(name, zone),
				AvailabilityZone:   zone,
				AvailabilityZoneID: fmt.Sprintf("%s-id", zone),
			})
			assert.Nil(t, natTask.Local.Create(context.Background(), &v1beta1.Subnet{
				ObjectMeta: metav1.ObjectMeta{
					Name: namer.PublicSubnetName(name, zone),
					Annotations: map[string]string{
						meta.AnnotationKeyExternalName: fmt.Sprintf("id-%s", zone),
					},
				},
			}))
		}
		sub, err := findPublicSubnetByZone(context.Background(), natTask.Local, natTask.VPC, zones[1])
		assert.Nil(t, err)
		assert.Equal(t, v1alpha1.FindSubnetByZone(natTask.VPC.Status.PublicSubnets, zones[1]).Name, sub.Name)
	})
}

func Test_findPrivateSubnetByZone(t *testing.T) {
	natTask := fakeNATTask(t)
	t.Run("subnet not found", func(t *testing.T) {
		sub, err := findPrivateSubnetByZone(context.Background(), natTask.Local, natTask.VPC, "zone")
		assert.Nil(t, sub)
		assert.ErrorContains(t, err, "no private subnet found")
	})
	t.Run("subnet found", func(t *testing.T) {
		zones := natTask.GetNodeZones()
		name := natTask.InternetName()
		for _, zone := range zones {
			natTask.VPC.Status.PrivateSubnets = append(natTask.VPC.Status.PrivateSubnets, v1alpha1.Subnet{
				Name:               namer.PrivateSubnetName(name, zone),
				AvailabilityZone:   zone,
				AvailabilityZoneID: fmt.Sprintf("%s-id", zone),
			})
			assert.Nil(t, natTask.Local.Create(context.Background(), &v1beta1.Subnet{
				ObjectMeta: metav1.ObjectMeta{
					Name: namer.PrivateSubnetName(name, zone),
					Annotations: map[string]string{
						meta.AnnotationKeyExternalName: fmt.Sprintf("id-%s", zone),
					},
				},
			}))
		}
		sub, err := findPrivateSubnetByZone(context.Background(), natTask.Local, natTask.VPC, zones[1])
		assert.Nil(t, err)
		assert.Equal(t, v1alpha1.FindSubnetByZone(natTask.VPC.Status.PrivateSubnets, zones[1]).Name, sub.Name)
	})
}

func Test_findNodeSubnetByZone(t *testing.T) {
	natTask := fakeNATTask(t)
	t.Run("subnet not found", func(t *testing.T) {
		sub, err := findNodeSubnetByZone(context.Background(), natTask.Local, natTask.VPC, "zone")
		assert.Nil(t, sub)
		assert.ErrorContains(t, err, "no node subnet found")
	})
	t.Run("subnet found", func(t *testing.T) {
		zones := natTask.GetNodeZones()
		name := natTask.InternetName()
		for _, zone := range zones {
			natTask.VPC.Status.NodeSubnets = append(natTask.VPC.Status.NodeSubnets, v1alpha1.Subnet{
				Name:               namer.NodesSubnetName(name, zone),
				AvailabilityZone:   zone,
				AvailabilityZoneID: fmt.Sprintf("%s-id", zone),
			})
			assert.Nil(t, natTask.Local.Create(context.Background(), &v1beta1.Subnet{
				ObjectMeta: metav1.ObjectMeta{
					Name: namer.NodesSubnetName(name, zone),
					Annotations: map[string]string{
						meta.AnnotationKeyExternalName: fmt.Sprintf("id-%s", zone),
					},
				},
			}))
		}
		sub, err := findNodeSubnetByZone(context.Background(), natTask.Local, natTask.VPC, zones[1])
		assert.Nil(t, err)
		assert.Equal(t, v1alpha1.FindSubnetByZone(natTask.VPC.Status.NodeSubnets, zones[1]).Name, sub.Name)
	})
}

func TestNAT_deleteNAT(t *testing.T) {
	natTask := fakeNATTask(t)
	t.Run("delete nat success", func(t *testing.T) {
		name := namer.PublicSubnetName("test", "zone")
		assert.Nil(t, natTask.Local.Create(context.Background(), &v1beta1.NATGateway{
			ObjectMeta: metav1.ObjectMeta{
				Name: name,
				Annotations: map[string]string{
					meta.AnnotationKeyExternalName: "test-id",
				},
			},
		}))
		assert.Nil(t, natTask.deleteNAT(context.Background(), "test", "zone"))
		assert.ErrorContains(t, natTask.Local.Get(context.Background(), types.NamespacedName{Name: name}, &v1beta1.NATGateway{}), "not found")
	})
	t.Run("delete nat failed", func(t *testing.T) {
		natTask := fakeNATTask(t)
		ctl := gomock.NewController(t)
		mockcli := kubemock.NewMockClient(ctl)
		natTask.Local = mockcli
		mockcli.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).Times(1)
		mockcli.EXPECT().Patch(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).Times(1)
		mockcli.EXPECT().Delete(gomock.Any(), gomock.Any()).Return(assert.AnError).Times(1)
		assert.ErrorIs(t, natTask.deleteNAT(context.Background(), "test", "zone"), assert.AnError)
	})
	t.Run("delete nat not found", func(t *testing.T) {
		natTask := fakeNATTask(t)
		ctl := gomock.NewController(t)
		mockcli := kubemock.NewMockClient(ctl)
		natTask.Local = mockcli
		mockcli.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(apierrors.NewNotFound(schema.GroupResource{}, "not found")).Times(1)
		mockcli.EXPECT().Patch(gomock.Any(), gomock.Any(), gomock.Any()).Return(apierrors.NewNotFound(schema.GroupResource{}, "not found")).Times(1)
		mockcli.EXPECT().Delete(gomock.Any(), gomock.Any()).Return(apierrors.NewNotFound(schema.GroupResource{}, "not found")).Times(1)
		assert.Nil(t, natTask.deleteNAT(context.Background(), "test", "zone"))
	})
}

func TestNAT_Create_TagsConsistencyAcrossZones(t *testing.T) {
	t.Run("tags should be consistent across all zones", func(t *testing.T) {
		natTask := fakeNATTask(t)

		// Set AdditionalCloudTags to trigger the bug in line 57
		natTask.Config.AdditionalCloudTags = map[string]string{
			"environment": "test",
			"team":        "infra",
		}

		zones := natTask.GetNodeZones()
		natTask.VPC.Status.PublicSubnets = []v1alpha1.Subnet{}
		natTask.Internet.Status.IPAddresses = []v1alpha1.IPAddress{}
		name := natTask.InternetName()

		// Setup test data for all zones
		for _, zone := range zones {
			natTask.VPC.Status.PublicSubnets = append(natTask.VPC.Status.PublicSubnets, v1alpha1.Subnet{
				Name:               namer.PublicSubnetName(name, zone),
				AvailabilityZone:   zone,
				AvailabilityZoneID: fmt.Sprintf("%s-id", zone),
			})
			assert.Nil(t, natTask.Local.Create(context.Background(), &v1beta1.Subnet{
				ObjectMeta: metav1.ObjectMeta{
					Name: namer.PublicSubnetName(name, zone),
					Annotations: map[string]string{
						meta.AnnotationKeyExternalName: fmt.Sprintf("id-%s", zone),
					},
				},
			}))
			natTask.Internet.Status.IPAddresses = append(natTask.Internet.Status.IPAddresses, v1alpha1.IPAddress{
				Name:         namer.PublicSubnetName(name, zone),
				AllocationID: fmt.Sprintf("allo-%s", zone),
			})
		}

		assert.Nil(t, natTask.Local.Status().Update(context.Background(), natTask.VPC))

		// Execute Create
		result := natTask.Create(context.Background())
		assert.True(t, result.IsDoneState())
		assert.Equal(t, 3, len(natTask.nats))

		// Verify tags consistency across all zones
		var allTags []map[string]string
		for i, zone := range zones {
			natRes := &v1beta1.NATGateway{}
			natName := namer.PublicSubnetName(name, zone)
			assert.Nil(t, natTask.Local.Get(context.Background(), types.NamespacedName{Name: natName}, natRes))

			// Get tags from the created NATGateway CR
			assert.NotNil(t, natRes.Spec.ForProvider.Tags, "Tags should not be nil for zone %s", zone)
			tagsWithPointers := natRes.Spec.ForProvider.Tags

			// Convert map[string]*string to map[string]string for easier comparison
			tags := make(map[string]string)
			for k, v := range tagsWithPointers {
				if v != nil {
					tags[k] = *v
				}
			}
			allTags = append(allTags, tags)

			// All tags should contain AdditionalCloudTags
			assert.Equal(t, "test", tags["environment"], "Zone %d should have environment tag", i)
			assert.Equal(t, "infra", tags["team"], "Zone %d should have team tag", i)

			// All tags should have Name field (but with different values per zone)
			assert.Contains(t, tags, "Name", "Zone %d should have Name tag", i)
		}

		// Compare tags between zones (excluding the "Name" field which should be different)
		zone1Tags := allTags[0]
		zone2Tags := allTags[1]
		zone3Tags := allTags[2]

		// Extract common tags (excluding "Name")
		zone1Common := make(map[string]string)
		zone2Common := make(map[string]string)
		zone3Common := make(map[string]string)

		for k, v := range zone1Tags {
			if k != "Name" {
				zone1Common[k] = v
			}
		}
		for k, v := range zone2Tags {
			if k != "Name" {
				zone2Common[k] = v
			}
		}
		for k, v := range zone3Tags {
			if k != "Name" {
				zone3Common[k] = v
			}
		}

		// All zones should have the same tags (except Name)
		assert.Equal(t, zone1Common, zone2Common, "Zone 1 and Zone 2 should have identical tags (except Name)")
		assert.Equal(t, zone1Common, zone3Common, "Zone 1 and Zone 3 should have identical tags (except Name)")
		assert.Equal(t, zone2Common, zone3Common, "Zone 2 and Zone 3 should have identical tags (except Name)")

		// Verify Name tags are different per zone
		assert.NotEqual(t, zone1Tags["Name"], zone2Tags["Name"], "Zone 1 and Zone 2 should have different Name tags")
		assert.NotEqual(t, zone1Tags["Name"], zone3Tags["Name"], "Zone 1 and Zone 3 should have different Name tags")
		assert.NotEqual(t, zone2Tags["Name"], zone3Tags["Name"], "Zone 2 and Zone 3 should have different Name tags")
	})
}

func TestNAT_Delete(t *testing.T) {
	natTask := fakeNATTask(t)
	natTask.Config.ResourceNamePrefix = "tidbcloud"
	zones := natTask.GetNodeZones()
	natTask.VPC.Status.PublicSubnets = []v1alpha1.Subnet{}
	natTask.Internet.Status.IPAddresses = []v1alpha1.IPAddress{}
	name := natTask.InternetName()
	for _, zone := range zones {
		natTask.VPC.Status.PublicSubnets = append(natTask.VPC.Status.PublicSubnets, v1alpha1.Subnet{
			Name:               namer.PublicSubnetName(name, zone),
			AvailabilityZone:   zone,
			AvailabilityZoneID: fmt.Sprintf("%s-id", zone),
		})
		assert.Nil(t, natTask.Local.Create(context.Background(), &v1beta1.Subnet{
			ObjectMeta: metav1.ObjectMeta{
				Name: namer.PublicSubnetName(name, zone),
				Annotations: map[string]string{
					meta.AnnotationKeyExternalName: fmt.Sprintf("id-%s", zone),
				},
			},
		}))
		natTask.Internet.Status.IPAddresses = append(natTask.Internet.Status.IPAddresses, v1alpha1.IPAddress{
			Name:         namer.PublicSubnetName(name, zone),
			AllocationID: fmt.Sprintf("allo-%s", zone),
		})
	}
	assert.Nil(t, natTask.Local.Status().Update(context.Background(), natTask.VPC))
	result := natTask.Create(context.Background())
	assert.True(t, result.IsDoneState())
	result = natTask.Delete(context.Background())
	assert.True(t, result.IsDoneState())
}