/*
  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

  Licensed under the Apache License, Version 2.0 (the "License").
  You may not use this file except in compliance with the License.
  You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
*/

package test

import (
	"context"
	"errors"
	"fmt"
	"log"
	"log/slog"
	"strconv"
	"testing"
	"time"

	"github.com/aws/aws-advanced-go-wrapper/.test/test_framework/container/test_utils"
	_ "github.com/aws/aws-advanced-go-wrapper/aws-secrets-manager"
	"github.com/aws/aws-advanced-go-wrapper/awssql/error_util"
	"github.com/aws/aws-advanced-go-wrapper/awssql/property_util"
	_ "github.com/aws/aws-advanced-go-wrapper/otlp"
	_ "github.com/aws/aws-advanced-go-wrapper/xray"
	"github.com/stretchr/testify/require"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/secretsmanager"
	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
)

func CreateSecret(t *testing.T, client *secretsmanager.Client, env *test_utils.TestEnvironment, secretName string) string {
	var engine string
	if env.Info().Request.Engine == test_utils.PG {
		engine = "postgres"
	} else if env.Info().Request.Engine == test_utils.MYSQL {
		engine = "mysql"
	}

	secretString := fmt.Sprintf(
		`{"engine":"%s","dbname":"%s","host":"%s","username":"%s","password":"%s","description":"Test secret generated by integration tests."}`,
		engine,
		env.Info().DatabaseInfo.DefaultDbName,
		env.Info().DatabaseInfo.ClusterEndpoint,
		env.Info().DatabaseInfo.Username,
		env.Info().DatabaseInfo.Password)

	createSecretInput := &secretsmanager.CreateSecretInput{
		Name:         aws.String(secretName),
		SecretString: aws.String(secretString),
	}
	createSecretOutput, err := client.CreateSecret(context.TODO(), createSecretInput)
	if err != nil {
		t.Fatalf("Failed to create secret: '%v'.", err)
	}
	return *createSecretOutput.ARN
}

func DeleteSecret(t *testing.T, client *secretsmanager.Client, secretName string) {
	deleteSecretInput := &secretsmanager.DeleteSecretInput{
		SecretId:                   aws.String(secretName),
		ForceDeleteWithoutRecovery: aws.Bool(true),
	}
	_, err := client.DeleteSecret(context.TODO(), deleteSecretInput)
	if err != nil {
		t.Fatalf("Failed to delete secret: '%v'.", err)
	}
}

func TestSecretsManager(t *testing.T) {
	env, err := test_utils.GetCurrentTestEnvironment()
	if env == nil || err != nil {
		t.Fail()
	}

	cfg, _ := config.LoadDefaultConfig(context.TODO())
	client := secretsmanager.NewFromConfig(cfg, func(o *secretsmanager.Options) {
		o.Region = env.Info().Region
	})
	secretName := fmt.Sprintf("TestSecret-%s", uuid.New().String())
	secretARN := CreateSecret(t, client, env, secretName)
	defer DeleteSecret(t, client, secretName)

	t.Run("Connection", func(t *testing.T) {
		connectionTest(t, env, secretName)
	})

	t.Run("ConnectionWithTelemetryOtel", func(t *testing.T) {
		connectionTestWithTelemetryOtel(t, env, secretName)
	})

	t.Run("ConnectionWithTelemetryXray", func(t *testing.T) {
		connectionTestWithTelemetryXray(t, env, secretName)
	})

	t.Run("ConnectWithARN", func(t *testing.T) {
		connectWithARNTest(t, env, secretARN)
	})

	t.Run("IncorrectSecretId", func(t *testing.T) {
		incorrectSecretIdTest(t, env)
	})

	t.Run("MissingSecretId", func(t *testing.T) {
		missingSecretIdTest(t, env)
	})

	t.Run("InvalidRegion", func(t *testing.T) {
		invalidRegionTest(t, env, secretName)
	})

	t.Run("MissingRegion", func(t *testing.T) {
		missingRegionTest(t, env, secretName)
	})

	t.Run("IncorrectRegion", func(t *testing.T) {
		incorrectRegionTest(t, env, secretName)
	})

	t.Run("WithFailover", func(t *testing.T) {
		failoverTest(t, env, secretName)
	})

	t.Run("WithFailoverEfm", func(t *testing.T) {
		failoverEfmTest(t, env, secretName)
	})

	t.Run("WithEfm", func(t *testing.T) {
		efmTest(t, env, secretName)
	})
}

func connectionTest(t *testing.T, env *test_utils.TestEnvironment, secretName string) {
	defer test_utils.BasicCleanupAfterBasicSetup(t)()
	environment, err := test_utils.GetCurrentTestEnvironment()
	assert.Nil(t, err)

	dsn := test_utils.GetDsn(env, map[string]string{
		"plugins":                "awsSecretsManager",
		"user":                   "incorrectUser",
		"password":               "incorrectPassword",
		"secretsManagerSecretId": secretName,
		"secretsManagerRegion":   env.Info().Region,
	})
	db, err := test_utils.OpenDb(environment.Info().Request.Engine, dsn)
	assert.Nil(t, err)
	assert.NotNil(t, db)
	defer db.Close()

	pingErr := db.Ping()
	assert.Nil(t, pingErr)
}

func connectionTestWithTelemetryOtel(t *testing.T, env *test_utils.TestEnvironment, secretName string) {
	defer test_utils.BasicCleanupAfterBasicSetup(t)()
	bsp, err := test_utils.SetupTelemetry(env)
	assert.Nil(t, err)
	assert.NotNil(t, bsp)
	defer func() { _ = bsp.Shutdown(context.TODO()) }()
	environment, err := test_utils.GetCurrentTestEnvironment()
	assert.Nil(t, err)

	dsn := test_utils.GetDsn(env, map[string]string{
		"plugins":                 "awsSecretsManager",
		"user":                    "incorrectUser",
		"password":                "incorrectPassword",
		"secretsManagerSecretId":  secretName,
		"secretsManagerRegion":    env.Info().Region,
		"enableTelemetry":         "true",
		"telemetryTracesBackend":  "OTLP",
		"telemetryMetricsBackend": "OTLP",
	})
	db, err := test_utils.OpenDb(environment.Info().Request.Engine, dsn)
	assert.Nil(t, err)
	assert.NotNil(t, db)
	defer db.Close()

	pingErr := db.Ping()
	assert.Nil(t, pingErr)
}

func connectionTestWithTelemetryXray(t *testing.T, env *test_utils.TestEnvironment, secretName string) {
	defer test_utils.BasicCleanupAfterBasicSetup(t)()
	bsp, err := test_utils.SetupTelemetry(env)
	assert.Nil(t, err)
	assert.NotNil(t, bsp)
	defer func() { _ = bsp.Shutdown(context.TODO()) }()
	environment, err := test_utils.GetCurrentTestEnvironment()
	assert.Nil(t, err)

	dsn := test_utils.GetDsn(env, map[string]string{
		"plugins":                 "awsSecretsManager",
		"user":                    "incorrectUser",
		"password":                "incorrectPassword",
		"secretsManagerSecretId":  secretName,
		"secretsManagerRegion":    env.Info().Region,
		"enableTelemetry":         "true",
		"telemetryTracesBackend":  "XRAY",
		"telemetryMetricsBackend": "OTLP",
	})
	db, err := test_utils.OpenDb(environment.Info().Request.Engine, dsn)
	assert.Nil(t, err)
	assert.NotNil(t, db)
	defer db.Close()

	pingErr := db.Ping()
	assert.Nil(t, pingErr)
}

func connectWithARNTest(t *testing.T, env *test_utils.TestEnvironment, secretARN string) {
	defer test_utils.BasicCleanupAfterBasicSetup(t)()
	environment, err := test_utils.GetCurrentTestEnvironment()
	assert.Nil(t, err)

	dsn := test_utils.GetDsn(env, map[string]string{
		"plugins":                "awsSecretsManager",
		"user":                   "incorrectUser",
		"password":               "incorrectPassword",
		"secretsManagerSecretId": secretARN,
	})
	db, err := test_utils.OpenDb(environment.Info().Request.Engine, dsn)
	assert.Nil(t, err)
	assert.NotNil(t, db)
	defer db.Close()

	pingErr := db.Ping()
	assert.Nil(t, pingErr)
}

func incorrectSecretIdTest(t *testing.T, env *test_utils.TestEnvironment) {
	defer test_utils.BasicCleanupAfterBasicSetup(t)()
	environment, err := test_utils.GetCurrentTestEnvironment()
	assert.Nil(t, err)

	dsn := test_utils.GetDsn(env, map[string]string{
		"plugins":                "awsSecretsManager",
		"user":                   "incorrectUser",
		"password":               "incorrectPassword",
		"secretsManagerSecretId": "incorrectSecretId",
		"secretsManagerRegion":   env.Info().Region,
	})
	db, err := test_utils.OpenDb(environment.Info().Request.Engine, dsn)
	assert.Nil(t, err)
	assert.NotNil(t, db)
	defer db.Close()

	pingErr := db.Ping()
	assert.NotNil(t, pingErr)
}

func missingSecretIdTest(t *testing.T, env *test_utils.TestEnvironment) {
	defer test_utils.BasicCleanupAfterBasicSetup(t)()
	environment, err := test_utils.GetCurrentTestEnvironment()
	assert.Nil(t, err)

	dsn := test_utils.GetDsn(env, map[string]string{
		"plugins":              "awsSecretsManager",
		"user":                 "incorrectUser",
		"password":             "incorrectPassword",
		"secretsManagerRegion": env.Info().Region,
	})
	db, err := test_utils.OpenDb(environment.Info().Request.Engine, dsn)
	assert.Nil(t, err)
	assert.NotNil(t, db)
	defer db.Close()

	pingErr := db.Ping()
	assert.NotNil(t, pingErr)
}

func invalidRegionTest(t *testing.T, env *test_utils.TestEnvironment, secretName string) {
	defer test_utils.BasicCleanupAfterBasicSetup(t)()
	environment, err := test_utils.GetCurrentTestEnvironment()
	assert.Nil(t, err)

	dsn := test_utils.GetDsn(env, map[string]string{
		"plugins":                "awsSecretsManager",
		"user":                   "incorrectUser",
		"password":               "incorrectPassword",
		"secretsManagerSecretId": secretName,
		"secretsManagerRegion":   "invalidRegion",
	})
	db, err := test_utils.OpenDb(environment.Info().Request.Engine, dsn)
	assert.Nil(t, err)
	assert.NotNil(t, db)
	defer db.Close()

	pingErr := db.Ping()
	assert.Equal(
		t,
		errors.New(error_util.GetMessage("AwsSecretsManagerConnectionPlugin.invalidRegion", "invalidRegion")),
		pingErr)
}

func missingRegionTest(t *testing.T, env *test_utils.TestEnvironment, secretName string) {
	defer test_utils.BasicCleanupAfterBasicSetup(t)()
	environment, err := test_utils.GetCurrentTestEnvironment()
	assert.Nil(t, err)

	dsn := test_utils.GetDsn(env, map[string]string{
		"plugins":                "awsSecretsManager",
		"user":                   "incorrectUser",
		"password":               "incorrectPassword",
		"secretsManagerSecretId": secretName,
	})
	db, err := test_utils.OpenDb(environment.Info().Request.Engine, dsn)
	assert.Nil(t, err)
	assert.NotNil(t, db)
	defer db.Close()

	pingErr := db.Ping()
	assert.Error(t, pingErr, "Default region does not match secret region, should throw an error")
}

func incorrectRegionTest(t *testing.T, env *test_utils.TestEnvironment, secretName string) {
	defer test_utils.BasicCleanupAfterBasicSetup(t)()
	environment, err := test_utils.GetCurrentTestEnvironment()
	assert.Nil(t, err)

	dsn := test_utils.GetDsn(env, map[string]string{
		"plugins":                "awsSecretsManager",
		"user":                   "incorrectUser",
		"password":               "incorrectPassword",
		"secretsManagerSecretId": secretName,
		"secretsManagerRegion":   "ca-central-1",
	})
	db, err := test_utils.OpenDb(environment.Info().Request.Engine, dsn)
	assert.Nil(t, err)
	assert.NotNil(t, db)
	defer db.Close()

	pingErr := db.Ping()
	assert.NotNil(t, pingErr)
}

func failoverTest(t *testing.T, env *test_utils.TestEnvironment, secretName string) {
	auroraTestUtility, environment, err := failoverSetup(t)
	defer test_utils.BasicCleanup(t.Name())
	assert.Nil(t, err)

	dsn := test_utils.GetDsn(env, map[string]string{
		"plugins":                "awsSecretsManager,failover",
		"user":                   "incorrectUser",
		"password":               "incorrectPassword",
		"secretsManagerSecretId": secretName,
		"secretsManagerRegion":   env.Info().Region,
	})
	wrapperDriver := test_utils.NewWrapperDriver(environment.Info().Request.Engine)

	conn, err := wrapperDriver.Open(dsn)
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close()

	// Check that we are connected to the writer.
	instanceId, err := test_utils.ExecuteInstanceQuery(environment.Info().Request.Engine, environment.Info().Request.Deployment, conn)
	assert.Nil(t, err)
	assert.True(t, auroraTestUtility.IsDbInstanceWriter(instanceId, ""))

	// Failover and check that it has failed over.
	triggerFailoverError := auroraTestUtility.FailoverClusterAndWaitTillWriterChanged(instanceId, "", "")
	assert.Nil(t, triggerFailoverError)
	_, queryError := test_utils.ExecuteInstanceQuery(environment.Info().Request.Engine, environment.Info().Request.Deployment, conn)
	require.Error(t, queryError, "Failover plugin did not complete failover successfully.")
	assert.Equal(t, error_util.GetMessage("Failover.connectionChangedError"), queryError.Error())

	// Assert that we are connected to the new writer after failover.
	newInstanceId, err := test_utils.ExecuteInstanceQuery(environment.Info().Request.Engine, environment.Info().Request.Deployment, conn)
	assert.Nil(t, err)
	currWriterId, err := auroraTestUtility.GetClusterWriterInstanceId("")
	assert.Nil(t, err)
	assert.Equal(t, currWriterId, newInstanceId)
	assert.NotEqual(t, instanceId, newInstanceId)

	test_utils.BasicCleanup(t.Name())
}

func efmTest(t *testing.T, env *test_utils.TestEnvironment, secretName string) {
	_, environment, err := failoverSetup(t)
	defer test_utils.BasicCleanup(t.Name())
	assert.Nil(t, err)
	props := test_utils.GetPropsForProxy(environment, environment.Info().ProxyDatabaseInfo.ClusterEndpoint, "awsSecretsManager,efm", TEST_FAILURE_DETECTION_INTERVAL_SECONDS)
	props[property_util.USER.Name] = "incorrectUser"
	props[property_util.PASSWORD.Name] = "incorrectPassword"
	props[property_util.SECRETS_MANAGER_REGION.Name] = env.Info().Region
	props[property_util.SECRETS_MANAGER_SECRET_ID.Name] = secretName
	props[property_util.FAILOVER_TIMEOUT_MS.Name] = strconv.Itoa(4 * TEST_MONITORING_TIMEOUT_SECONDS * 1000)

	dsn := test_utils.GetDsn(environment, props)
	db, err := test_utils.OpenDb(environment.Info().Request.Engine, dsn)
	assert.Nil(t, err)
	assert.NotNil(t, db)
	defer db.Close()

	// Verify connection works.
	err = db.Ping()
	require.NoError(t, err, "Failed to open database connection.")

	// Start a long-running query in a goroutine
	queryChan := make(chan error)
	go func() {
		// Execute a sleep query that will run for 10 seconds
		sleepQuery := test_utils.GetSleepSql(environment.Info().Request.Engine, TEST_SLEEP_QUERY_SECONDS)

		_, err := test_utils.ExecuteQueryDB(environment.Info().Request.Engine, db, sleepQuery, TEST_SLEEP_QUERY_TIMEOUT_SECONDS)
		queryChan <- err
	}()

	// Wait a bit to ensure the query has started
	time.Sleep(5 * time.Second)

	// Disable connectivity while the sleep query is running
	slog.Debug("Disabling all connectivity.")
	test_utils.DisableAllConnectivity()

	// Wait for the query to complete and check the error
	queryErr := <-queryChan
	close(queryChan)
	require.NotNil(t, queryErr)
	slog.Debug(fmt.Sprintf("Sleep query fails with error: %s.", queryErr.Error()))
	assert.False(t, errors.Is(queryErr, context.DeadlineExceeded), "Sleep query should have failed due to connectivity loss")

	// Re-enable connectivity
	slog.Debug("Re-enabling all connectivity.")
	test_utils.EnableAllConnectivity(true)

	newInstanceId, err := test_utils.ExecuteInstanceQueryDB(environment.Info().Request.Engine, environment.Info().Request.Deployment, db)
	assert.Nil(t, err, "After connectivity is re-enabled new connections should not throw errors.")
	assert.NotZero(t, newInstanceId)
}

func failoverEfmTest(t *testing.T, env *test_utils.TestEnvironment, secretName string) {
	auroraTestUtility, environment, err := failoverSetup(t)
	defer test_utils.BasicCleanup(t.Name())
	assert.Nil(t, err)
	props := test_utils.GetPropsForProxy(
		environment,
		environment.Info().ProxyDatabaseInfo.ClusterEndpoint,
		"awsSecretsManager,failover,efm",
		TEST_FAILURE_DETECTION_INTERVAL_SECONDS,
	)
	props[property_util.USER.Name] = "incorrectUser"
	props[property_util.PASSWORD.Name] = "incorrectPassword"
	props[property_util.SECRETS_MANAGER_REGION.Name] = env.Info().Region
	props[property_util.SECRETS_MANAGER_SECRET_ID.Name] = secretName
	props[property_util.FAILOVER_TIMEOUT_MS.Name] = strconv.Itoa(4 * TEST_MONITORING_TIMEOUT_SECONDS * 1000)

	dsn := test_utils.GetDsn(environment, props)
	db, err := test_utils.OpenDb(environment.Info().Request.Engine, dsn)
	assert.Nil(t, err)
	assert.NotNil(t, db)
	defer db.Close()

	// Verify connection works.
	err = db.Ping()
	require.NoError(t, err, "Failed to open database connection.")

	// Start a long-running query in a goroutine
	queryChan := make(chan error)
	go func() {
		// Execute a sleep query that will run for 10 seconds
		sleepQuery := test_utils.GetSleepSql(environment.Info().Request.Engine, TEST_SLEEP_QUERY_SECONDS)

		_, err := test_utils.ExecuteQueryDB(environment.Info().Request.Engine, db, sleepQuery, TEST_SLEEP_QUERY_TIMEOUT_SECONDS)
		queryChan <- err
	}()

	// Wait a bit to ensure the query has started
	time.Sleep(5 * time.Second)

	// Disable connectivity while the sleep query is running
	slog.Debug("Disabling all connectivity.")
	test_utils.DisableAllConnectivity()

	// Wait a bit to ensure failover has started
	time.Sleep(2 * TEST_MONITORING_TIMEOUT_SECONDS * time.Second)

	// Re-enable connectivity
	slog.Debug("Re-enabling all connectivity.")
	test_utils.EnableAllConnectivity(true)

	// Wait for the query to complete and check the error
	queryErr := <-queryChan
	close(queryChan)
	require.NotNil(t, queryErr)
	assert.Equal(t, error_util.GetMessage("Failover.connectionChangedError"), queryErr.Error())

	newInstanceId, err := test_utils.ExecuteInstanceQueryDB(environment.Info().Request.Engine, environment.Info().Request.Deployment, db)
	require.True(t, auroraTestUtility.IsDbInstanceWriter(newInstanceId, ""))
	assert.Nil(t, err, "After connectivity is re-enabled new connections should not throw errors.")
	assert.NotZero(t, newInstanceId)
}
