package aws

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"

	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	"github.com/stretchr/testify/assert"
	ec2v1beta1 "github.com/upbound/provider-aws/apis/ec2/v1beta1"
	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"
)

func makeKubernetesAgent() *agentv2.KubernetesAgent {
	return &agentv2.KubernetesAgent{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "test-agent",
			Namespace: "default",
			Annotations: map[string]string{
				consts.AnnotationKeyO11yEndpointServiceName: "epsrv-mock",
				consts.AnnotationKeyO11yServiceDNSName:      "test.service.com",
			},
			Labels: map[string]string{
				agentv2.TagKeyTenant: "tenant-1",
			},
		},
		Spec: agentv2.KubernetesAgentSpec{
			KubernetesClusterName: "cluster-1",
			EKS: &agentv2.EKSConfig{
				VpcID:           "vpc-123",
				AssumeRoleArn:   "arn:aws:iam::123:role/role",
				SecurityGroupID: "sg-123",
				NodeSubnetIDs:   []string{"subnet-1"},
			},
		},
	}
}

type runTestOpts struct {
	handler        http.HandlerFunc
	mutateAgent    func(*agentv2.KubernetesAgent)
	setupResources func(client.Client)
	serverless     bool
	errMsg         string
}

func runTest(opts runTestOpts, testFn func(*reconciler.Context[*agentv2.KubernetesAgent], *agentv2.KubernetesAgent)) {
	var serverURL string
	if opts.handler != nil {
		ts := httptest.NewServer(opts.handler)
		defer ts.Close()
		serverURL = ts.URL
	}

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

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

	cfg := &config.RootConfig{
		O11yConfig: config.ObservabilityInfra{
			Version: 1,
			Region:  "us-east-1",
		},
	}
	if opts.serverless {
		cfg.Mode = config.ModeServerless
	}

	ctx := &reconciler.Context[*agentv2.KubernetesAgent]{
		Agent:  agent,
		Local:  fakeCli,
		AWS:    testkits.NewAwsClient(serverURL),
		Config: cfg,
	}

	testFn(ctx, agent)
}

const (
	awsDescribeVpcEndpointsEmpty    = `<DescribeVpcEndpointsResponse xmlns="http://ec2.amazonaws.com/doc/2016-11-15/"><requestId>req-1</requestId><vpcEndpointSet/></DescribeVpcEndpointsResponse>`
	awsDescribeVpcEndpointsExisting = `<DescribeVpcEndpointsResponse xmlns="http://ec2.amazonaws.com/doc/2016-11-15/"><requestId>req-1</requestId><vpcEndpointSet><item><vpcEndpointId>vpce-existing</vpcEndpointId><serviceName>epsrv-mock</serviceName><vpcId>vpc-123</vpcId><vpcEndpointType>Interface</vpcEndpointType><state>Available</state></item></vpcEndpointSet></DescribeVpcEndpointsResponse>`
	awsDescribeVpcEndpointServices  = `<DescribeVpcEndpointServicesResponse xmlns="http://ec2.amazonaws.com/doc/2016-11-15/"><requestId>req-2</requestId><serviceDetailSet><item><serviceName>epsrv-mock</serviceName><availabilityZoneSet><item>us-east-1a</item></availabilityZoneSet></item></serviceDetailSet></DescribeVpcEndpointServicesResponse>`
	awsDescribeSubnets              = `<DescribeSubnetsResponse xmlns="http://ec2.amazonaws.com/doc/2016-11-15/"><requestId>req-3</requestId><subnetSet><item><subnetId>subnet-1</subnetId><availabilityZone>us-east-1a</availabilityZone></item></subnetSet></DescribeSubnetsResponse>`
	awsCreateVpcEndpointSuccess     = `<CreateVpcEndpointResponse xmlns="http://ec2.amazonaws.com/doc/2016-11-15/"><requestId>req-4</requestId><vpcEndpoint><vpcEndpointId>vpce-123456</vpcEndpointId><state>Available</state></vpcEndpoint></CreateVpcEndpointResponse>`
	awsDeleteVpcEndpointsSuccess    = `<DeleteVpcEndpointsResponse xmlns="http://ec2.amazonaws.com/doc/2016-11-15/"><requestId>req-1</requestId><unsuccessful/></DeleteVpcEndpointsResponse>`
)

func makeHandler(responses map[string]string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		_ = r.ParseForm()
		action := r.Form.Get("Action")
		if response, ok := responses[action]; ok {
			w.Header().Set("Content-Type", "application/xml")
			w.Write([]byte(response))
		} else {
			w.WriteHeader(http.StatusBadRequest)
		}
	}
}

func TestPrivateLinkTask_Create(t *testing.T) {
	tests := []struct {
		name           string
		serverless     bool
		handler        http.HandlerFunc
		mutateAgent    func(*agentv2.KubernetesAgent)
		setupResources func(client.Client)
		errMsg         string
		verify         func(*testing.T, client.Client, *agentv2.KubernetesAgent)
	}{
		{
			name:       "serverless: success create new endpoint",
			serverless: true,
			handler: makeHandler(map[string]string{
				"DescribeVpcEndpoints":        awsDescribeVpcEndpointsEmpty,
				"DescribeVpcEndpointServices": awsDescribeVpcEndpointServices,
				"DescribeSubnets":             awsDescribeSubnets,
				"CreateVpcEndpoint":           awsCreateVpcEndpointSuccess,
			}),
			verify: func(t *testing.T, _ client.Client, a *agentv2.KubernetesAgent) {
				endpointID := util.GetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName)
				assert.Equal(t, "vpce-123456", endpointID)
			},
		},
		{
			name:       "serverless: reuse existing endpoint",
			serverless: true,
			handler: makeHandler(map[string]string{
				"DescribeVpcEndpoints": awsDescribeVpcEndpointsExisting,
			}),
			verify: func(t *testing.T, _ client.Client, a *agentv2.KubernetesAgent) {
				endpointID := util.GetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName)
				assert.Equal(t, "vpce-existing", endpointID)
			},
		},
		{
			name:       "serverless: skip if endpoint already exists and version not changed",
			serverless: true,
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				util.SetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName, "vpce-existing")
				util.SetAnnotation(a, consts.AnnotationKeyO11yVersion, "1")
			},
			verify: func(t *testing.T, _ client.Client, a *agentv2.KubernetesAgent) {
				assert.Equal(t, "vpce-existing", util.GetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName))
			},
		},
		{
			name:       "serverless: endpoint service name not found",
			serverless: true,
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				delete(a.Annotations, consts.AnnotationKeyO11yEndpointServiceName)
			},
			errMsg: "endpoint service name not found",
		},
		{
			name:       "serverless: service dns name not found",
			serverless: true,
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				delete(a.Annotations, consts.AnnotationKeyO11yServiceDNSName)
			},
			handler: makeHandler(map[string]string{
				"DescribeVpcEndpoints": awsDescribeVpcEndpointsEmpty,
			}),
			errMsg: "service dns name not found",
		},
		{
			name:       "serverless: failed to describe vpc endpoints",
			serverless: true,
			handler: func(w http.ResponseWriter, r *http.Request) {
				w.WriteHeader(http.StatusInternalServerError)
			},
			errMsg: "failed to find reusable vpc endpoint",
		},
		{
			name:       "serverless: failed to create vpc endpoint",
			serverless: true,
			handler: makeHandler(map[string]string{
				"DescribeVpcEndpoints":        awsDescribeVpcEndpointsEmpty,
				"DescribeVpcEndpointServices": awsDescribeVpcEndpointServices,
				"DescribeSubnets":             awsDescribeSubnets,
				"CreateVpcEndpoint":           "",
			}),
			errMsg: "failed to create interface endpoint",
		},
		{
			name: "upbound: success create new endpoint",
			handler: makeHandler(map[string]string{
				"DescribeVpcEndpointServices": awsDescribeVpcEndpointServices,
				"DescribeSubnets":             awsDescribeSubnets,
			}),
			verify: func(t *testing.T, c client.Client, a *agentv2.KubernetesAgent) {
				epName := util.GetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName)
				assert.NotEmpty(t, epName)

				var ep ec2v1beta1.VPCEndpoint
				err := c.Get(context.Background(), client.ObjectKey{Name: epName}, &ep)
				assert.NoError(t, err)
				assert.Equal(t, "epsrv-mock", *ep.Spec.ForProvider.ServiceName)

				var subnetAssoc ec2v1beta1.VPCEndpointSubnetAssociation
				err = c.Get(context.Background(), client.ObjectKey{Name: fmt.Sprintf("%s-subnet-1", epName)}, &subnetAssoc)
				assert.NoError(t, err)
				assert.Equal(t, "subnet-1", *subnetAssoc.Spec.ForProvider.SubnetID)
			},
		},
		{
			name: "upbound: skip if endpoint already exists and version not changed",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				util.SetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName, "existing-ep")
				util.SetAnnotation(a, consts.AnnotationKeyO11yVersion, "1")
			},
			verify: func(t *testing.T, _ client.Client, a *agentv2.KubernetesAgent) {
				assert.Equal(t, "existing-ep", util.GetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName))
			},
		},
		{
			name: "upbound: reuse existing endpoint",
			setupResources: func(c client.Client) {
				reusableEP := &ec2v1beta1.VPCEndpoint{ObjectMeta: metav1.ObjectMeta{Name: "reusable-ep"}}
				reusableEP.Spec.ForProvider.ServiceName = ptr.To("epsrv-mock")
				reusableEP.Spec.ForProvider.VPCID = ptr.To("vpc-123")
				reusableEP.Spec.ForProvider.VPCEndpointType = ptr.To("Interface")
				reusableEP.Status.SetConditions(xpv1.Available())
				assert.NoError(t, c.Create(context.Background(), reusableEP))
			},
			verify: func(t *testing.T, _ client.Client, a *agentv2.KubernetesAgent) {
				assert.Equal(t, "reusable-ep", util.GetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName))
			},
		},
		{
			name: "upbound: endpoint service name not found",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				delete(a.Annotations, consts.AnnotationKeyO11yEndpointServiceName)
			},
			errMsg: "endpoint service name not found",
		},
		{
			name: "upbound: service dns name not found",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				delete(a.Annotations, consts.AnnotationKeyO11yServiceDNSName)
			},
			errMsg: "service dns name not found",
		},
		{
			name: "upbound: failed to create interface endpoint - describe error",
			handler: func(w http.ResponseWriter, r *http.Request) {
				w.WriteHeader(http.StatusInternalServerError)
			},
			errMsg: "failed to create interface endpoint",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			opts := runTestOpts{
				handler:        tt.handler,
				mutateAgent:    tt.mutateAgent,
				setupResources: tt.setupResources,
				serverless:     tt.serverless,
				errMsg:         tt.errMsg,
			}
			runTest(opts, func(ctx *reconciler.Context[*agentv2.KubernetesAgent], agent *agentv2.KubernetesAgent) {
				task := NewPrivateLinkTask(ctx)
				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, agent)
				}
			})
		})
	}
}

func TestPrivateLinkTask_Delete(t *testing.T) {
	tests := []struct {
		name           string
		serverless     bool
		handler        http.HandlerFunc
		mutateAgent    func(*agentv2.KubernetesAgent)
		setupResources func(client.Client)
		errMsg         string
		verify         func(*testing.T, client.Client, *agentv2.KubernetesAgent)
	}{
		{
			name:       "serverless: success delete endpoint",
			serverless: true,
			handler:    makeHandler(map[string]string{"DeleteVpcEndpoints": awsDeleteVpcEndpointsSuccess}),
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				util.SetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName, "vpce-123456")
			},
			verify: func(t *testing.T, _ client.Client, a *agentv2.KubernetesAgent) {
				annotation := util.GetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName)
				assert.Empty(t, annotation)
			},
		},
		{
			name:       "serverless: skip deletion if endpoint id is empty",
			serverless: true,
		},
		{
			name:       "serverless: failed to delete vpc endpoint",
			serverless: true,
			handler: func(w http.ResponseWriter, r *http.Request) {
				w.WriteHeader(http.StatusInternalServerError)
			},
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				util.SetAnnotation(a, consts.AnnotationKeyO11ySvcInterfaceEndpointName, "vpce-123456")
			},
			errMsg: "failed to delete vpc interface endpoint",
		},
		{
			name: "upbound: success delete",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				// Keep default dnsName from makeKubernetesAgent
			},
			setupResources: func(c client.Client) {
				agent := makeKubernetesAgent()
				epName := util.InterfaceEndpointName(agent.ObjectMeta, "test.service.com", consts.O11yEndpointNamePrefixDataplaneNextGen)
				ep := &ec2v1beta1.VPCEndpoint{
					ObjectMeta: metav1.ObjectMeta{Name: epName},
				}
				subnetAssoc := &ec2v1beta1.VPCEndpointSubnetAssociation{
					ObjectMeta: metav1.ObjectMeta{Name: fmt.Sprintf("%s-subnet-1", epName)},
				}
				sgAssoc := &ec2v1beta1.VPCEndpointSecurityGroupAssociation{
					ObjectMeta: metav1.ObjectMeta{Name: fmt.Sprintf("%s-sg-123", epName)},
				}

				assert.NoError(t, c.Create(context.Background(), ep))
				assert.NoError(t, c.Create(context.Background(), subnetAssoc))
				assert.NoError(t, c.Create(context.Background(), sgAssoc))
			},
			verify: func(t *testing.T, c client.Client, a *agentv2.KubernetesAgent) {
				epName := util.InterfaceEndpointName(a.ObjectMeta, "test.service.com", consts.O11yEndpointNamePrefixDataplaneNextGen)

				var ep ec2v1beta1.VPCEndpoint
				err := c.Get(context.Background(), client.ObjectKey{Name: epName}, &ep)
				assert.True(t, apierrors.IsNotFound(err))

				var subnetAssoc ec2v1beta1.VPCEndpointSubnetAssociation
				err = c.Get(context.Background(), client.ObjectKey{Name: fmt.Sprintf("%s-subnet-1", epName)}, &subnetAssoc)
				assert.True(t, apierrors.IsNotFound(err))

				var sgAssoc ec2v1beta1.VPCEndpointSecurityGroupAssociation
				err = c.Get(context.Background(), client.ObjectKey{Name: fmt.Sprintf("%s-sg-123", epName)}, &sgAssoc)
				assert.True(t, apierrors.IsNotFound(err))
			},
		},
		{
			name: "upbound: resources already deleted",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				// Keep default dnsName from makeKubernetesAgent, but don't setup resources
			},
		},
		{
			name: "upbound: no dns name configured",
			mutateAgent: func(a *agentv2.KubernetesAgent) {
				delete(a.Annotations, consts.AnnotationKeyO11yServiceDNSName)
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			opts := runTestOpts{
				handler:        tt.handler,
				mutateAgent:    tt.mutateAgent,
				setupResources: tt.setupResources,
				serverless:     tt.serverless,
				errMsg:         tt.errMsg,
			}
			runTest(opts, func(ctx *reconciler.Context[*agentv2.KubernetesAgent], agent *agentv2.KubernetesAgent) {
				task := NewPrivateLinkTask(ctx)
				res := task.Delete(context.Background())
				if tt.errMsg != "" {
					assert.Error(t, res.Error())
					assert.Contains(t, res.Error().Error(), tt.errMsg)
				} else {
					assert.Equal(t, deploy.Done(), res)
				}

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