package cluster

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"

	"github.com/tidbcloud/eks-provider/tests/e2e/dag"
	"github.com/tidbcloud/eks-provider/tests/e2e/env"
	"github.com/tidbcloud/eks-provider/tests/e2e/label"
	"github.com/tidbcloud/eks-provider/tests/e2e/suite"
	"github.com/tidbcloud/eks-provider/tests/e2e/utils/helper"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/cluster-controller/tests/cmd/end-user-proxy/mysql"
)

func init() {
	suite.AddCase(
		label.New[dag.Case](
			&IsRunning{},
			"Cluster should be running",
		),
		label.New[dag.Case](
			&PublicEndpointIsConnectable{},
			"cluster's public endpoint should be connectable",
			label.FeatPublicEndpoint,
		),
		label.New(
			NewPrivateEndpointIsConnectable(helper.ConnectionTypePeering),
			"cluster's peering endpoint should be connectable",
		),
		label.New(
			NewPrivateEndpointIsConnectable(helper.ConnectionTypePrivateLink),
			"cluster's private link should be connectable",
		),
		label.New[dag.Case](
			&PublicConnIsAlwaysActive{},
			"cluster's public endpoint should not be down",
			label.FeatPublicEndpoint,
		),
		label.New(
			NewPrivateConnIsAlwaysActive(helper.ConnectionTypePeering),
			"cluster's peering endpoint should not be down",
		),
		label.New(
			NewPrivateConnIsAlwaysActive(helper.ConnectionTypePrivateLink),
			"cluster's private link endpoint should not be down",
		),
	)
}

type IsRunning struct {
	Cluster *cloudv1.Cluster
}

func (c *IsRunning) Test(t *testing.T) {
	assert.Equal(t, cloudv1.ClusterPhaseRunning, c.Cluster.Status.Phase)
}

func (c *IsRunning) Envs() []string {
	return []string{
		env.WithCloudAndVersion("Cluster", env.CloudAWS, env.InfraProviderVersionOld),
		env.WithCloudAndVersion("Cluster", env.CloudAWS, env.InfraProviderVersionNew),
		// env.WithCloudAndVersion("Cluster", env.CloudGCP, env.InfraProviderVersionOld),
		// env.WithCloudAndVersion("Cluster", env.CloudGCP, env.InfraProviderVersionNew),
	}
}

type PublicEndpointIsConnectable struct {
	Cluster *cloudv1.Cluster
}

func (c *PublicEndpointIsConnectable) Test(t *testing.T) {
	dsn := helper.GetDSN(&c.Cluster.Status.SqlEndpoints, helper.ConnectionTypePublic)
	assert.NotEmpty(t, dsn)
	ac := mysql.Action{
		Command:        mysql.CommandPing,
		DataSourceName: dsn,
	}
	err := mysql.NewLocal().Invoke(context.Background(), &ac)
	assert.NoError(t, err)

	for i, group := range c.Cluster.Spec.ExtraGroups {
		status := c.Cluster.Status.ExtraEndpoints[i]
		assert.Equal(t, status.Name, group.Name)
		dsn := helper.GetDSN(&status.SqlEndpointsStatus, helper.ConnectionTypePublic)
		assert.NotEmpty(t, dsn)
		ac := mysql.Action{
			Command:        mysql.CommandPing,
			DataSourceName: dsn,
		}
		err := mysql.NewLocal().Invoke(context.Background(), &ac)
		assert.NoError(t, err)
	}
}

func (c *PublicEndpointIsConnectable) Envs() []string {
	return []string{
		env.WithCloudAndVersion("Cluster", env.CloudAWS, env.InfraProviderVersionOld),
		env.WithCloudAndVersion("Cluster", env.CloudAWS, env.InfraProviderVersionNew),
		// env.WithCloudAndVersion("Cluster", env.CloudGCP, env.InfraProviderVersionOld),
		// env.WithCloudAndVersion("Cluster", env.CloudGCP, env.InfraProviderVersionNew),

		// endpoint is connectable after upgrading
		env.WithCloud("UpgradedCluster", env.CloudAWS),

		// all public endpoint of extra groups should work
		env.WithCloudAndVersion("UpdateCluster:ExtraGroup", env.CloudAWS, env.InfraProviderVersionNew),
	}
}

type PrivateEndpointIsConnectable struct {
	Config  *env.Config
	Cluster *cloudv1.Cluster

	AWSEndUserLambdaFn string

	connType helper.ConnectionType
}

func NewPrivateEndpointIsConnectable(connType helper.ConnectionType) dag.Case {
	return &PrivateEndpointIsConnectable{
		connType: connType,
	}
}

func (c *PrivateEndpointIsConnectable) Test(t *testing.T) {
	dsn := helper.GetDSN(&c.Cluster.Status.SqlEndpoints, c.connType)
	assert.NotEmpty(t, dsn)
	ac := mysql.Action{
		Command:        mysql.CommandPing,
		DataSourceName: dsn,
	}
	err := mysql.NewLambda(c.Config.AWSConfig, c.AWSEndUserLambdaFn).Invoke(context.Background(), &ac)
	assert.NoError(t, err)
}

func (c *PrivateEndpointIsConnectable) Envs() []string {
	return []string{
		env.WithCloudAndVersion("ClusterWithEndUser", env.CloudAWS, env.InfraProviderVersionOld),
		env.WithCloudAndVersion("ClusterWithEndUser", env.CloudAWS, env.InfraProviderVersionNew),
	}
}

type PublicConnIsAlwaysActive struct {
	Cluster *cloudv1.Cluster
}

func (c *PublicConnIsAlwaysActive) Test(t *testing.T) {
	dsn := helper.GetDSN(&c.Cluster.Status.SqlEndpoints, helper.ConnectionTypePublic)
	assert.NotEmpty(t, dsn)

	var err error
	ac := mysql.Action{
		Command:        mysql.CommandCheckConn,
		DataSourceName: dsn,
	}
	go func() {
		ctx := context.Background()
		err = mysql.NewLocal().Invoke(ctx, &ac)
	}()
	t.Parallel()
	// check error at the end of test
	assert.NoError(t, err, "public connection should not be down")
}

func (c *PublicConnIsAlwaysActive) Envs() []string {
	return []string{
		// connection should keep alive for default cluster with old infra provider
		env.WithCloudAndVersion("Cluster", env.CloudAWS, env.InfraProviderVersionOld),
		// env.WithCloudAndVersion("Cluster", env.CloudAWS, env.InfraProviderVersionNew),
		// env.WithCloudAndVersion("Cluster", env.CloudGCP, env.InfraProviderVersionOld),
		// env.WithCloudAndVersion("Cluster", env.CloudGCP, env.InfraProviderVersionNew),
	}
}

type PrivateConnIsAlwaysActive struct {
	Config  *env.Config
	Cluster *cloudv1.Cluster

	AWSEndUserLambdaFn string

	connType helper.ConnectionType
}

func NewPrivateConnIsAlwaysActive(connType helper.ConnectionType) dag.Case {
	return &PrivateConnIsAlwaysActive{
		connType: connType,
	}
}

func (c *PrivateConnIsAlwaysActive) Test(t *testing.T) {
	dsn := helper.GetDSN(&c.Cluster.Status.SqlEndpoints, c.connType)
	assert.NotEmpty(t, dsn)

	var err error
	ac := mysql.Action{
		Command:        mysql.CommandCheckConn,
		DataSourceName: dsn,
		Timeout:        time.Minute * 10,
	}
	go func() {
		ctx := context.Background()
		err = mysql.NewLambda(c.Config.AWSConfig, c.AWSEndUserLambdaFn).Invoke(ctx, &ac)
	}()
	t.Parallel()
	// check error at the end of test
	assert.NoError(t, err, "private connection should not be down")
}

func (c *PrivateConnIsAlwaysActive) Envs() []string {
	return []string{
		env.WithCloudAndVersion("ClusterWithEndUser", env.CloudAWS, env.InfraProviderVersionOld),
		env.WithCloudAndVersion("ClusterWithEndUser", env.CloudAWS, env.InfraProviderVersionNew),
	}
}