package integration

import (
	"context"
	"encoding/base64"
	"fmt"
	"net"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	v1cluster "edge-infra.dev/pkg/edge/apis/cluster/v1alpha1"
	"edge-infra.dev/pkg/sds/remoteaccess/constants"
	v1vpnconfig "edge-infra.dev/pkg/sds/remoteaccess/k8s/apis/vpnconfigs/v1"
	r "edge-infra.dev/pkg/sds/remoteaccess/wireguard/relay"
	"edge-infra.dev/pkg/sds/remoteaccess/wireguard/store"
	"edge-infra.dev/test/f2"
	"edge-infra.dev/test/f2/x/ktest"
)

var f f2.Framework

var (
	projectID     = "ret-edge-b79we3ikmc7j9mihuwst2"
	testPublicKey = "wbWuHrJEPC2Ui7XVQoWuM/8HZAG1FlLC/08L2vvEEgw="

	outOfSubnetIPAddress = "179.16.17.3"
	inSubnetIPAddress    = "172.16.16.12"

	clusterAName = "cluster-a"
	clusterBName = "cluster-b"
	clusterCName = "cluster-c"
	testClusterA = createCluster(clusterAName)
	testClusterB = createCluster(clusterBName)
	testClusterC = createCluster(clusterCName)

	testVPNConfig                     = createVPNConfig(clusterAName, inSubnetIPAddress)
	testVPNConfigNoIPAddress          = createVPNConfig(clusterCName, "")
	testVPNConfigOutOfSubnetIPAddress = createVPNConfig(clusterBName, outOfSubnetIPAddress)

	expectedRelaySecretInterfaceLines = []string{
		"[Interface]",
		"ListenPort = 51820",
		"PostUp = iptables -A FORWARD -i %i -j ACCEPT; iptables -A FORWARD -o %i -j ACCEPT; iptables -t nat -A POSTROUTING -o eth0 -j MASQUERADE",
		"PostDown = iptables -D FORWARD -i %i -j ACCEPT; iptables -D FORWARD -o %i -j ACCEPT; iptables -t nat -D POSTROUTING -o eth0 -j MASQUERADE",
	}
	expectedRelaySecretPeerLines = []string{
		"",
		"[Peer]",
	}

	vpnNamespace = &corev1.Namespace{
		ObjectMeta: metav1.ObjectMeta{
			Name: constants.VPNNamespace,
		},
	}
)

func TestMain(m *testing.M) {
	f = f2.New(context.Background(),
		f2.WithExtensions(
			ktest.New(),
		)).
		Setup(func(ctx f2.Context) (f2.Context, error) {
			k, err := ktest.FromContext(ctx)
			if err != nil {
				return ctx, err
			}
			// Override timeouts if we aren't using a live cluster
			if !*k.Env.UseExistingCluster {
				k.Timeout = 5 * time.Second
				k.Tick = 10 * time.Millisecond
			}
			return ctx, nil
		}).Teardown()
	os.Exit(f.Run(m))
}

func TestRelayWireguard(t *testing.T) {
	var (
		relay    *r.Relay
		storeAWg *store.Store
		storeBWg *store.Store
		storeCWg *store.Store
		subnet   *net.IPNet
		clientIP net.IP
	)
	feature := f2.NewFeature("RelayWireguard").
		Setup("create wireguard instances", func(ctx f2.Context, t *testing.T) f2.Context {
			var err error
			k := ktest.FromContextT(ctx, t)
			require.NoError(t, k.Client.Create(ctx, vpnNamespace.DeepCopy()))

			relay, err = r.Get(ctx, k.Client)
			require.NoError(t, err)

			storeAWg, err = store.Get(ctx, k.Client, testVPNConfigNoIPAddress.DeepCopy(), testClusterA)
			require.NoError(t, err)

			storeBWg, err = store.Get(ctx, k.Client, testVPNConfigOutOfSubnetIPAddress.DeepCopy(), testClusterB)
			require.NoError(t, err)

			storeCWg, err = store.Get(ctx, k.Client, testVPNConfig.DeepCopy(), testClusterC)
			require.NoError(t, err)
			return ctx
		}).
		Setup("create subnet and client IP", func(ctx f2.Context, t *testing.T) f2.Context {
			var err error
			_, subnet, err = net.ParseCIDR("172.16.16.0/28")
			require.NoError(t, err)

			clientIP = net.ParseIP("172.16.16.1")
			return ctx
		}).
		Test("wireguard relay secret contains expected data", func(ctx f2.Context, t *testing.T) f2.Context {
			storeConfigs := map[string]*store.Store{clusterAName: storeAWg}
			secret := relay.GenerateConfigurationSecret(subnet, clientIP, testPublicKey, storeConfigs)
			secretLines := getLinesFromWireguardSecretData(secret)

			// assert expected lines appear in interface and peer sections
			assert.Equal(t, expectedRelaySecretInterfaceLines, secretLines[:4])
			assert.Equal(t, expectedRelaySecretPeerLines, secretLines[7:9])

			// check relay interface address is set
			expectedInterfaceLine := fmt.Sprintf("Address = %s", subnet)
			assert.Equal(t, expectedInterfaceLine, secretLines[4])

			// private key is not exposed, so check string is present and base64 encoded
			assert.Equal(t, "PrivateKey = ", secretLines[5][:13])
			privatekey := secretLines[5][13:]
			_, err := base64.StdEncoding.DecodeString(privatekey)
			assert.NoError(t, err)

			// check relay interface MTU is set
			expectedMTULine := fmt.Sprintf("MTU = %s", constants.MTU)
			assert.Equal(t, expectedMTULine, secretLines[6])

			// check friendly name is set for peer config
			expectedClientFriendlyName := "# friendly_json={\"cluster_name\":\"cluster_infra\"}"
			assert.Equal(t, expectedClientFriendlyName, secretLines[9])

			// check client public key is set in peer config
			expectedClientPublicKeyLine := fmt.Sprintf("PublicKey = %s", testPublicKey)
			assert.Equal(t, expectedClientPublicKeyLine, secretLines[11])

			// check store friendly name is set in peer config
			expectedStoreFriendlyName := fmt.Sprintf("# friendly_json={\"cluster_name\":\"%s\",\"cluster\":\"%s\",\"vpn_enabled\":\"%t\"}", "4c4d-30-05-22", clusterAName, true)
			assert.Equal(t, expectedStoreFriendlyName, secretLines[14])

			// check store ip address is set in peer config
			expectedStoreAllowedIPsLine := fmt.Sprintf("AllowedIPs = %s/32", storeConfigs[clusterAName].GetIPAddress())
			assert.Equal(t, expectedStoreAllowedIPsLine, secretLines[15])

			// check store public key is set in peer config
			expectedStorePublicKeyline := fmt.Sprintf("PublicKey = %s", storeConfigs[clusterAName].GetPublicKey())
			assert.Equal(t, expectedStorePublicKeyline, secretLines[16])
			return ctx
		}).
		Test("multiple stores are added to relay secret", func(ctx f2.Context, t *testing.T) f2.Context {
			// add store one and check peer is added to relay config
			storeConfigs := map[string]*store.Store{clusterAName: storeAWg}
			secret := relay.GenerateConfigurationSecret(subnet, clientIP, testPublicKey, storeConfigs)
			secretLines := getLinesFromWireguardSecretData(secret)
			assert.Len(t, secretLines, 18)
			assert.Equal(t, "[Peer]", secretLines[8])

			// add store two and check another peer is added to relay config
			storeConfigs = map[string]*store.Store{clusterAName: storeAWg, clusterBName: storeBWg}
			secret = relay.GenerateConfigurationSecret(subnet, clientIP, testPublicKey, storeConfigs)
			secretLines = getLinesFromWireguardSecretData(secret)
			assert.Len(t, secretLines, 23)
			assert.Equal(t, "[Peer]", secretLines[13])

			// add store three and check another peer is added to relay config
			storeConfigs = map[string]*store.Store{clusterAName: storeAWg, clusterBName: storeBWg, clusterCName: storeCWg}
			secret = relay.GenerateConfigurationSecret(subnet, clientIP, testPublicKey, storeConfigs)
			secretLines = getLinesFromWireguardSecretData(secret)
			assert.Len(t, secretLines, 28)
			assert.Equal(t, "[Peer]", secretLines[18])
			return ctx
		}).
		Test("store can be removed from relay secret", func(ctx f2.Context, t *testing.T) f2.Context {
			// add store and check peer is added to relay config
			storeConfigs := map[string]*store.Store{clusterAName: storeAWg}
			secret := relay.GenerateConfigurationSecret(subnet, clientIP, testPublicKey, storeConfigs)
			secretLines := getLinesFromWireguardSecretData(secret)
			assert.Len(t, secretLines, 18)
			assert.Equal(t, "[Peer]", secretLines[8])

			// remove store and check peer is removed from relay config
			storeConfigs = map[string]*store.Store{}
			secret = relay.GenerateConfigurationSecret(subnet, clientIP, testPublicKey, storeConfigs)
			secretLines = getLinesFromWireguardSecretData(secret)
			assert.Len(t, secretLines, 13)
			return ctx
		}).Feature()
	f.Test(t, feature)
}

func getLinesFromWireguardSecretData(secret *corev1.Secret) []string {
	secretData := secret.StringData[constants.WireguardSecretField]
	return strings.Split(secretData, "\n")
}

func createCluster(name string) *v1cluster.Cluster {
	return &v1cluster.Cluster{
		ObjectMeta: metav1.ObjectMeta{Name: name},
		Spec: v1cluster.ClusterSpec{
			Banner:       "dev0-zynstra",
			Fleet:        "store",
			Location:     "us-east1-c",
			Name:         "4c4d-30-05-22",
			Organization: "edge-dev0-retail-gmi062",
			ProjectID:    projectID,
			Type:         "sds",
		},
	}
}

func createVPNConfig(name, ip string) *v1vpnconfig.VPNConfig {
	return &v1vpnconfig.VPNConfig{
		TypeMeta:   metav1.TypeMeta{Kind: "VPNConfig", APIVersion: "remoteaccess.edge.ncr.com"},
		ObjectMeta: metav1.ObjectMeta{Namespace: constants.VPNNamespace, Name: name, UID: "1234"},
		Spec: v1vpnconfig.VPNConfigSpec{
			Enabled: true,
		},
		Status: &v1vpnconfig.VPNConfigStatus{
			IP: ip,
		},
	}
}