package commonazure

import (
	"context"
	"testing"

	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	"github.com/stretchr/testify/assert"
	azurenetworkv1beta2 "github.com/upbound/provider-azure/apis/network/v1beta2"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"

	agentv2 "github.com/tidbcloud/observability/agent-control-server/api/v2"
	"github.com/tidbcloud/observability/agent-control-server/pkg/config"
	"github.com/tidbcloud/observability/agent-control-server/pkg/consts"
	"github.com/tidbcloud/observability/agent-control-server/pkg/deploy"
	"github.com/tidbcloud/observability/agent-control-server/pkg/reconciler"
	"github.com/tidbcloud/observability/agent-control-server/pkg/testkits"
	"github.com/tidbcloud/observability/agent-control-server/pkg/util"
	"github.com/tidbcloud/observability/agent-control-server/pkg/util/namer"
)

func makeTestAgent() *agentv2.KubernetesAgent {
	return &agentv2.KubernetesAgent{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "test-agent",
			Namespace: "default",
			Annotations: map[string]string{
				consts.AnnotationKeyO11yEndpointServiceName:  "test-service-alias",
				consts.AnnotationKeyO11yServiceOriginDNSName: "test.service.com",
			},
			Labels: map[string]string{agentv2.TagKeyTenant: "tenant-1"},
		},
		Spec: agentv2.KubernetesAgentSpec{
			KubernetesClusterName: "cluster-1",
			AKS: &agentv2.AKSConfig{
				SubnetID:      "/subscriptions/sub-1/resourceGroups/rg-1/providers/Microsoft.Network/virtualNetworks/vnet-1/subnets/subnet-1",
				ResourceGroup: "rg-1",
			},
		},
	}
}

func testEndpointName(agent *agentv2.KubernetesAgent) string {
	return namer.UniqueIdentifier(agent.ObjectMeta, namer.WithPrefix(consts.O11yEndpointNamePrefixDataplaneNextGen))
}

func newTestPrivateLinkTask(ctx *reconciler.Context[*agentv2.KubernetesAgent]) *PrivateLinkTask[*agentv2.KubernetesAgent] {
	return NewPrivateLinkTask(ctx, PrivateLinkConfig{
		AKS:                   func() *agentv2.AKSConfig { return ctx.Agent.Spec.AKS },
		KubernetesClusterName: func() string { return ctx.Agent.Spec.KubernetesClusterName },
		ObjectMeta:            func() metav1.ObjectMeta { return ctx.Agent.ObjectMeta },
		ClusterMetaFunc: func() deploy.ClusterMeta {
			return deploy.ClusterMeta{
				K8sClusterName: ctx.Agent.Spec.KubernetesClusterName,
			}
		},
	})
}

type privateLinkTestOpts struct {
	mutateAgent    func(*agentv2.KubernetesAgent)
	setupResources func(client.Client)
	region         string
	errMsg         string
}

func runPrivateLinkTest(opts privateLinkTestOpts, testFn func(*reconciler.Context[*agentv2.KubernetesAgent], *agentv2.KubernetesAgent)) {
	scheme := testkits.Scheme()
	if err := azurenetworkv1beta2.SchemeBuilder.AddToScheme(scheme); err != nil {
		panic(err)
	}

	agent := makeTestAgent()
	if opts.mutateAgent != nil {
		opts.mutateAgent(agent)
	}

	fakeCli := testkits.BuildClient(agent)
	if opts.setupResources != nil {
		opts.setupResources(fakeCli)
	}

	region := opts.region
	if region == "" {
		region = "eastus"
	}

	ctx := &reconciler.Context[*agentv2.KubernetesAgent]{
		Agent: agent,
		Local: fakeCli,
		Config: &config.RootConfig{
			O11yConfig: config.ObservabilityInfra{
				Version: 1,
				Region:  region,
			},
		},
	}

	testFn(ctx, agent)
}

func TestPrivateLinkTask_Create(t *testing.T) {
	tests := []struct {
		name           string
		mutateAgent    func(*agentv2.KubernetesAgent)
		setupResources func(client.Client)
		errMsg         string
		verify         func(*testing.T, client.Client, *agentv2.KubernetesAgent)
	}{
		{
			name: "success: create new endpoint",
			verify: func(t *testing.T, c client.Client, a *agentv2.KubernetesAgent) {
				epName := util.GetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName)
				assert.NotEmpty(t, epName)

				var ep azurenetworkv1beta2.PrivateEndpoint
				err := c.Get(context.Background(), client.ObjectKey{Name: epName}, &ep)
				assert.NoError(t, err)
				assert.Equal(t, "test-service-alias", *ep.Spec.InitProvider.PrivateServiceConnection.PrivateConnectionResourceAlias)
				assert.Equal(t, a.Spec.AKS.SubnetID, *ep.Spec.ForProvider.SubnetID)
			},
		},
		{
			name: "success: create without backend dns name",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				delete(a.Annotations, consts.AnnotationKeyO11yServiceOriginDNSName)
			},
			verify: func(t *testing.T, c client.Client, a *agentv2.KubernetesAgent) {
				epName := util.GetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName)
				assert.NotEmpty(t, epName)

				var ep azurenetworkv1beta2.PrivateEndpoint
				err := c.Get(context.Background(), client.ObjectKey{Name: epName}, &ep)
				assert.NoError(t, err)
				assert.Equal(t, epName, testEndpointName(a))
			},
		},
		{
			name: "skip if endpoint exists and version unchanged",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				util.SetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName, "existing-pe")
				util.SetAnnotation(a, consts.AnnotationKeyO11yVersion, "1")
			},
			verify: func(t *testing.T, _ client.Client, a *agentv2.KubernetesAgent) {
				assert.Equal(t, "existing-pe", util.GetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName))
			},
		},
		{
			name: "endpoint service name not found",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				delete(a.Annotations, consts.AnnotationKeyO11yEndpointServiceName)
			},
			errMsg: "endpoint service name not found",
		},
		{
			name: "reuse existing endpoint",
			setupResources: func(c client.Client) {
				ep := &azurenetworkv1beta2.PrivateEndpoint{
					ObjectMeta: metav1.ObjectMeta{Name: "reusable-pe"},
					Spec: azurenetworkv1beta2.PrivateEndpointSpec{
						ForProvider: azurenetworkv1beta2.PrivateEndpointParameters{
							PrivateServiceConnection: &azurenetworkv1beta2.PrivateServiceConnectionParameters{
								PrivateConnectionResourceAlias: ptr.To("test-service-alias"),
							},
							SubnetID: ptr.To("/subscriptions/sub-1/resourceGroups/rg-1/providers/Microsoft.Network/virtualNetworks/vnet-1/subnets/subnet-1"),
						},
					},
				}
				ep.Status.SetConditions(xpv1.Available())
				assert.NoError(t, c.Create(context.Background(), ep))
			},
			verify: func(t *testing.T, _ client.Client, a *agentv2.KubernetesAgent) {
				assert.Equal(t, "reusable-pe", util.GetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName))
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			runPrivateLinkTest(privateLinkTestOpts{
				mutateAgent:    tt.mutateAgent,
				setupResources: tt.setupResources,
				region:         "eastus",
				errMsg:         tt.errMsg,
			}, func(ctx *reconciler.Context[*agentv2.KubernetesAgent], a *agentv2.KubernetesAgent) {
				task := newTestPrivateLinkTask(ctx)
				assert.Equal(t, "PrivateLink", task.Name())
				assert.Equal(t, "cluster-1", task.ClusterMeta().K8sClusterName)

				res := task.Create(context.Background())
				if tt.errMsg != "" {
					assert.Contains(t, res.Error().Error(), tt.errMsg)
				} else {
					assert.Equal(t, deploy.Done(), res)
				}

				if tt.verify != nil {
					tt.verify(t, ctx.Local, a)
				}
			})
		})
	}
}

func TestPrivateLinkTask_PostCreate(t *testing.T) {
	tests := []struct {
		name           string
		mutateAgent    func(*agentv2.KubernetesAgent)
		setupResources func(client.Client)
		errMsg         string
		verify         func(*testing.T, *agentv2.KubernetesAgent)
	}{
		{
			name: "success: ready with CustomDNSConfigs IP",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				util.SetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName, "pe-ready-dns")
			},
			setupResources: func(c client.Client) {
				ep := &azurenetworkv1beta2.PrivateEndpoint{ObjectMeta: metav1.ObjectMeta{Name: "pe-ready-dns"}}
				ep.Status.SetConditions(xpv1.Available())
				ep.Status.AtProvider.CustomDNSConfigs = []azurenetworkv1beta2.CustomDNSConfigsObservation{{
					IPAddresses: []*string{ptr.To("10.0.0.1")},
				}}
				assert.NoError(t, c.Create(context.Background(), ep))
			},
			verify: func(t *testing.T, a *agentv2.KubernetesAgent) {
				assert.Equal(t, "10.0.0.1", util.GetAnnotation(a, consts.AnnotationKeyO11yServiceDNSName))
			},
		},
		{
			name: "success: ready with PrivateServiceConnection IP",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				util.SetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName, "pe-ready-psc")
			},
			setupResources: func(c client.Client) {
				ep := &azurenetworkv1beta2.PrivateEndpoint{ObjectMeta: metav1.ObjectMeta{Name: "pe-ready-psc"}}
				ep.Status.SetConditions(xpv1.Available())
				ep.Status.AtProvider.PrivateServiceConnection = &azurenetworkv1beta2.PrivateServiceConnectionObservation{
					PrivateIPAddress: ptr.To("10.0.0.2"),
				}
				assert.NoError(t, c.Create(context.Background(), ep))
			},
			verify: func(t *testing.T, a *agentv2.KubernetesAgent) {
				assert.Equal(t, "10.0.0.2", util.GetAnnotation(a, consts.AnnotationKeyO11yServiceDNSName))
			},
		},
		{
			name: "PE not found",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				util.SetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName, "pe-missing")
			},
			errMsg: "failed to get private endpoint",
		},
		{
			name: "PE not ready",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				util.SetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName, "pe-not-ready")
			},
			setupResources: func(c client.Client) {
				ep := &azurenetworkv1beta2.PrivateEndpoint{ObjectMeta: metav1.ObjectMeta{Name: "pe-not-ready"}}
				ep.Status.SetConditions(xpv1.Creating())
				assert.NoError(t, c.Create(context.Background(), ep))
			},
			errMsg: "private endpoint is not ready",
		},
		{
			name: "PE ready but no IP",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				util.SetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName, "pe-no-ip")
			},
			setupResources: func(c client.Client) {
				ep := &azurenetworkv1beta2.PrivateEndpoint{ObjectMeta: metav1.ObjectMeta{Name: "pe-no-ip"}}
				ep.Status.SetConditions(xpv1.Available())
				assert.NoError(t, c.Create(context.Background(), ep))
			},
			errMsg: "private endpoint IP not found",
		},
		{
			name:   "no endpoint name annotation",
			verify: nil,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			runPrivateLinkTest(privateLinkTestOpts{
				mutateAgent:    tt.mutateAgent,
				setupResources: tt.setupResources,
				region:         "eastus",
				errMsg:         tt.errMsg,
			}, func(ctx *reconciler.Context[*agentv2.KubernetesAgent], a *agentv2.KubernetesAgent) {
				task := newTestPrivateLinkTask(ctx)
				res := task.PostCreate(context.Background())
				if tt.errMsg != "" {
					assert.Contains(t, res.Error().Error(), tt.errMsg)
				} else {
					assert.Equal(t, deploy.Done(), res)
				}

				if tt.verify != nil {
					tt.verify(t, a)
				}
			})
		})
	}
}

func TestPrivateLinkTask_Delete(t *testing.T) {
	tests := []struct {
		name           string
		mutateAgent    func(*agentv2.KubernetesAgent)
		setupResources func(client.Client)
		errMsg         string
		verify         func(*testing.T, client.Client)
	}{
		{
			name: "success: PE deleted from fake client",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				util.SetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName, "pe-delete")
			},
			setupResources: func(c client.Client) {
				ep := &azurenetworkv1beta2.PrivateEndpoint{ObjectMeta: metav1.ObjectMeta{Name: "pe-delete"}}
				assert.NoError(t, c.Create(context.Background(), ep))
			},
			verify: func(t *testing.T, c client.Client) {
				var ep azurenetworkv1beta2.PrivateEndpoint
				err := c.Get(context.Background(), client.ObjectKey{Name: "pe-delete"}, &ep)
				assert.True(t, apierrors.IsNotFound(err))
			},
		},
		{
			name: "PE already gone",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				util.SetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName, "pe-missing")
			},
		},
		{
			name: "delete derived endpoint without annotations",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				delete(a.Annotations, consts.AnnotationKeyO11ySvcInterfaceEndpointName)
				delete(a.Annotations, consts.AnnotationKeyO11yServiceOriginDNSName)
			},
			setupResources: func(c client.Client) {
				ep := &azurenetworkv1beta2.PrivateEndpoint{ObjectMeta: metav1.ObjectMeta{Name: testEndpointName(makeTestAgent())}}
				assert.NoError(t, c.Create(context.Background(), ep))
			},
			verify: func(t *testing.T, c client.Client) {
				var ep azurenetworkv1beta2.PrivateEndpoint
				err := c.Get(context.Background(), client.ObjectKey{Name: testEndpointName(makeTestAgent())}, &ep)
				assert.True(t, apierrors.IsNotFound(err))
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			runPrivateLinkTest(privateLinkTestOpts{
				mutateAgent:    tt.mutateAgent,
				setupResources: tt.setupResources,
				region:         "eastus",
				errMsg:         tt.errMsg,
			}, func(ctx *reconciler.Context[*agentv2.KubernetesAgent], _ *agentv2.KubernetesAgent) {
				task := newTestPrivateLinkTask(ctx)
				res := task.Delete(context.Background())
				if tt.errMsg != "" {
					assert.Contains(t, res.Error().Error(), tt.errMsg)
				} else {
					assert.Equal(t, deploy.Done(), res)
				}

				if tt.verify != nil {
					tt.verify(t, ctx.Local)
				}
			})
		})
	}
}