package jumper

import (
	"context"
	"time"

	"github.com/go-logr/logr"
	"github.com/onsi/ginkgo/v2"
	"github.com/onsi/gomega"
	jmpapi "github.com/tidbcloud/devops/api/jump"
	jpent "github.com/tidbcloud/devops/entity/jump"
	k8sutil "github.com/tidbcloud/devops/pkg/infra/utils/k8s"
	"github.com/tidbcloud/devops/pkg/service/jump/jumppod/types"
	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes/fake"
	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
	"k8s.io/client-go/tools/clientcmd"
	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
	ctrlclient "sigs.k8s.io/controller-runtime/pkg/client"
	ctrlfake "sigs.k8s.io/controller-runtime/pkg/client/fake"
)

var _ = ginkgo.Describe("[K8sTokenJumper RBAC]", func() {
	ctx := context.Background()

	ginkgo.It("creates primary RBAC resources when gateway access is disabled", func() {
		j, jctx := newTestK8sTokenJumper("test-user")
		dest := &jpent.DestinationControlPlane{Name: "infra-apiserver"}
		param := &types.PrepareJumpParam{Privilege: jpent.JumpPrivilegeReadWrite}

		err := j.prepareRBAC(ctx, jctx, dest, param)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		_, err = j.cli.Kube.CoreV1().Namespaces().Get(ctx, jctx.Namer.Namespace, metav1.GetOptions{})
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		jumpSA := &corev1.ServiceAccount{}
		err = j.cli.Ctrl.Get(ctx, ctrlclient.ObjectKey{
			Name:      jctx.Namer.JumpServiceAccount().Name,
			Namespace: jctx.Namer.JumpServiceAccount().Namespace,
		}, jumpSA)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		jumpRoleBinding := &rbacv1.ClusterRoleBinding{}
		err = j.cli.Ctrl.Get(ctx, ctrlclient.ObjectKey{Name: jctx.Namer.JumpClusterRoleBinding().Name}, jumpRoleBinding)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		gomega.Expect(jumpRoleBinding.RoleRef.Name).To(gomega.Equal(JumpReadWriteRoleName))
		gomega.Expect(jumpRoleBinding.Subjects).To(gomega.HaveLen(1))
		gomega.Expect(jumpRoleBinding.Subjects[0].Name).To(gomega.Equal(jctx.Namer.JumpServiceAccount().Name))
		gomega.Expect(jumpRoleBinding.Subjects[0].Namespace).To(gomega.Equal(jctx.Namer.JumpServiceAccount().Namespace))

		gatewaySA := &corev1.ServiceAccount{}
		err = j.cli.Ctrl.Get(ctx, ctrlclient.ObjectKey{
			Name:      jctx.Namer.JumpGatewayServiceAccount().Name,
			Namespace: jctx.Namer.JumpGatewayServiceAccount().Namespace,
		}, gatewaySA)
		gomega.Expect(apierrors.IsNotFound(err)).To(gomega.BeTrue())
	})

	ginkgo.It("creates gateway RBAC resources when gateway access is enabled", func() {
		j, jctx := newTestK8sTokenJumper("test-user")
		dest := &jpent.DestinationControlPlane{Name: "infra-apiserver"}
		param := &types.PrepareJumpParam{
			Privilege:           jpent.JumpPrivilegeAdmin,
			EnableGatewayAccess: true,
		}

		err := j.prepareRBAC(ctx, jctx, dest, param)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		gatewaySA := &corev1.ServiceAccount{}
		err = j.cli.Ctrl.Get(ctx, ctrlclient.ObjectKey{
			Name:      jctx.Namer.JumpGatewayServiceAccount().Name,
			Namespace: jctx.Namer.JumpGatewayServiceAccount().Namespace,
		}, gatewaySA)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		gatewayRole := &rbacv1.ClusterRole{}
		err = j.cli.Ctrl.Get(ctx, ctrlclient.ObjectKey{Name: jctx.Namer.JumpGatewayClusterRole().Name}, gatewayRole)
		gomega.Expect(apierrors.IsNotFound(err)).To(gomega.BeTrue())

		gatewayRoleBinding := &rbacv1.ClusterRoleBinding{}
		err = j.cli.Ctrl.Get(ctx, ctrlclient.ObjectKey{Name: jctx.Namer.JumpGatewayClusterRoleBinding().Name}, gatewayRoleBinding)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		gomega.Expect(gatewayRoleBinding.RoleRef.Name).To(gomega.Equal(JumpGatewayReadOnlyRoleName))
		gomega.Expect(gatewayRoleBinding.Subjects).To(gomega.HaveLen(1))
		gomega.Expect(gatewayRoleBinding.Subjects[0].Name).To(gomega.Equal(jctx.Namer.JumpGatewayServiceAccount().Name))
		gomega.Expect(gatewayRoleBinding.Subjects[0].Namespace).To(gomega.Equal(jctx.Namer.JumpGatewayServiceAccount().Namespace))
	})

	ginkgo.It("cleanup removes both primary and gateway resources", func() {
		j, jctx := newTestK8sTokenJumper("test-user")
		dest := &jpent.DestinationControlPlane{Name: "infra-apiserver"}
		param := &types.PrepareJumpParam{
			Privilege:           jpent.JumpPrivilegeReadWrite,
			EnableGatewayAccess: true,
		}

		err := j.prepareRBAC(ctx, jctx, dest, param)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		err = j.cleanupRBAC(ctx, jctx, dest)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		expectNotFound := func(obj ctrlclient.Object, key ctrlclient.ObjectKey) {
			err := j.cli.Ctrl.Get(ctx, key, obj)
			gomega.Expect(apierrors.IsNotFound(err)).To(gomega.BeTrue())
		}

		expectNotFound(&corev1.ServiceAccount{}, ctrlclient.ObjectKey{
			Name:      jctx.Namer.JumpServiceAccount().Name,
			Namespace: jctx.Namer.JumpServiceAccount().Namespace,
		})
		expectNotFound(&rbacv1.ClusterRole{}, ctrlclient.ObjectKey{Name: jctx.Namer.JumpClusterRole().Name})
		expectNotFound(&rbacv1.ClusterRoleBinding{}, ctrlclient.ObjectKey{Name: jctx.Namer.JumpClusterRoleBinding().Name})

		expectNotFound(&corev1.ServiceAccount{}, ctrlclient.ObjectKey{
			Name:      jctx.Namer.JumpGatewayServiceAccount().Name,
			Namespace: jctx.Namer.JumpGatewayServiceAccount().Namespace,
		})
		expectNotFound(&rbacv1.ClusterRole{}, ctrlclient.ObjectKey{Name: jctx.Namer.JumpGatewayClusterRole().Name})
		expectNotFound(&rbacv1.ClusterRoleBinding{}, ctrlclient.ObjectKey{Name: jctx.Namer.JumpGatewayClusterRoleBinding().Name})
	})

	ginkgo.It("cleanup keeps resources when destination does not match annotation", func() {
		j, jctx := newTestK8sTokenJumper("test-user")
		createdDest := &jpent.DestinationControlPlane{Name: "infra-apiserver"}
		param := &types.PrepareJumpParam{
			Privilege:           jpent.JumpPrivilegeReadWrite,
			EnableGatewayAccess: true,
		}

		err := j.prepareRBAC(ctx, jctx, createdDest, param)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		otherDest := &jpent.DestinationControlPlane{Name: "another-control-plane"}
		err = j.cleanupRBAC(ctx, jctx, otherDest)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		err = j.cli.Ctrl.Get(ctx, ctrlclient.ObjectKey{
			Name:      jctx.Namer.JumpServiceAccount().Name,
			Namespace: jctx.Namer.JumpServiceAccount().Namespace,
		}, &corev1.ServiceAccount{})
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		err = j.cli.Ctrl.Get(ctx, ctrlclient.ObjectKey{
			Name:      jctx.Namer.JumpGatewayServiceAccount().Name,
			Namespace: jctx.Namer.JumpGatewayServiceAccount().Namespace,
		}, &corev1.ServiceAccount{})
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
	})
})

func newTestK8sTokenJumper(jumpName string) (*K8sTokenJumper, *types.JumpContext) {
	scheme := runtime.NewScheme()
	err := clientgoscheme.AddToScheme(scheme)
	gomega.Expect(err).ToNot(gomega.HaveOccurred())

	ctrlCli := ctrlfake.NewClientBuilder().WithScheme(scheme).Build()
	kubeCli := fake.NewSimpleClientset()

	j := &K8sTokenJumper{
		cli: &k8sutil.UnionClient{
			Ctrl: ctrlCli,
			Kube: kubeCli,
		},
		roleCreator: fakeJumpRoleCreator{},
	}

	id := types.NewID(jumpName)
	jctx := &types.JumpContext{
		Name:   jumpName,
		ID:     id,
		Namer:  types.NewResourceNamer(id, "jump-system"),
		Logger: logr.Discard(),
	}
	return j, jctx
}

var _ = ginkgo.Describe("[K8sTokenJumper Token/Cred]", func() {
	ctx := context.Background()

	ginkgo.It("stores only primary cred when gateway access is disabled", func() {
		j, jctx, kubeCli := newTestK8sTokenJumperWithConfig("test-user")
		jp := &fakeJumpPodForJumper{jctx: jctx}
		seedServiceAccountToken(kubeCli, jctx.Namer.JumpServiceAccount().Namespace, jctx.Namer.JumpServiceAccount().Name, "jump-secret", "jump-token")

		dest := &jpent.DestinationControlPlane{Name: "infra-apiserver"}
		param := types.GenerateTokenParam{ExpireAt: time.Now().Add(time.Hour)}
		jmpParam, err := j.prepareToken(ctx, jp, dest, param, false, K8sJumpOption{})
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		credID := string(jpent.FormatCredentialID(dest))
		gomega.Expect(jp.storedKubeCred).To(gomega.HaveLen(1))
		gomega.Expect(jp.storedKubeCred).To(gomega.HaveKey(credID))
		gomega.Expect(jp.storedKubeCred).ToNot(gomega.HaveKey(string(jpent.FormatGatewayCredentialID(dest))))
		gomega.Expect(jmpParam.Attr.K8s.EnableGatewayAccess).To(gomega.BeFalse())

		cfg, err := clientcmd.Load(jp.storedKubeCred[credID])
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		gomega.Expect(cfg.CurrentContext).To(gomega.Equal(credID))
	})

	ginkgo.It("stores primary and gateway cred when gateway access is enabled", func() {
		j, jctx, kubeCli := newTestK8sTokenJumperWithConfig("test-user")
		jp := &fakeJumpPodForJumper{jctx: jctx}
		seedServiceAccountToken(kubeCli, jctx.Namer.JumpServiceAccount().Namespace, jctx.Namer.JumpServiceAccount().Name, "jump-secret", "jump-token")
		seedServiceAccountToken(kubeCli, jctx.Namer.JumpGatewayServiceAccount().Namespace, jctx.Namer.JumpGatewayServiceAccount().Name, "gateway-secret", "gateway-token")

		dest := &jpent.DestinationControlPlane{Name: "infra-apiserver"}
		param := types.GenerateTokenParam{ExpireAt: time.Now().Add(time.Hour)}
		jmpParam, err := j.prepareToken(ctx, jp, dest, param, true, K8sJumpOption{})
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		credID := string(jpent.FormatCredentialID(dest))
		gatewayCredID := string(jpent.FormatGatewayCredentialID(dest))
		gomega.Expect(jp.storedKubeCred).To(gomega.HaveLen(2))
		gomega.Expect(jp.storedKubeCred).To(gomega.HaveKey(credID))
		gomega.Expect(jp.storedKubeCred).To(gomega.HaveKey(gatewayCredID))
		gomega.Expect(jmpParam.Attr.K8s.EnableGatewayAccess).To(gomega.BeTrue())

		primaryCfg, err := clientcmd.Load(jp.storedKubeCred[credID])
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		gomega.Expect(primaryCfg.CurrentContext).To(gomega.Equal(credID))
		gatewayCfg, err := clientcmd.Load(jp.storedKubeCred[gatewayCredID])
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		gomega.Expect(gatewayCfg.CurrentContext).To(gomega.Equal(gatewayCredID))
	})

	ginkgo.It("refreshes both creds when jump attr enables gateway access", func() {
		j, jctx, kubeCli := newTestK8sTokenJumperWithConfig("test-user")
		jp := &fakeJumpPodForJumper{jctx: jctx}
		seedServiceAccountToken(kubeCli, jctx.Namer.JumpServiceAccount().Namespace, jctx.Namer.JumpServiceAccount().Name, "jump-secret", "jump-token")
		seedServiceAccountToken(kubeCli, jctx.Namer.JumpGatewayServiceAccount().Namespace, jctx.Namer.JumpGatewayServiceAccount().Name, "gateway-secret", "gateway-token")

		dest := &jpent.DestinationControlPlane{Name: "infra-apiserver"}
		attr := jpent.JumpAttributes{
			K8s: &jpent.JumpK8sAttributes{
				EnableGatewayAccess: true,
			},
		}
		param := &types.RefreshJumpParam{
			GenerateTokenParam: types.GenerateTokenParam{
				ExpireAt: time.Now().Add(time.Hour),
			},
		}
		_, err := j.RefreshJump(ctx, jp, dest, attr, param, K8sJumpOption{})
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		gomega.Expect(jp.storedKubeCred).To(gomega.HaveLen(2))
		gomega.Expect(jp.storedKubeCred).To(gomega.HaveKey(string(jpent.FormatCredentialID(dest))))
		gomega.Expect(jp.storedKubeCred).To(gomega.HaveKey(string(jpent.FormatGatewayCredentialID(dest))))
	})

	ginkgo.It("finalize removes both primary and gateway creds", func() {
		j, jctx := newTestK8sTokenJumper("test-user")
		jp := &fakeJumpPodForJumper{jctx: jctx}
		dest := &jpent.DestinationControlPlane{Name: "infra-apiserver"}

		err := j.FinalizeJump(ctx, jp, dest)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		gomega.Expect(jp.removedKubeCred).To(gomega.Equal([]string{
			string(jpent.FormatCredentialID(dest)),
			string(jpent.FormatGatewayCredentialID(dest)),
		}))
	})
})

type fakeJumpPodForJumper struct {
	jctx            *types.JumpContext
	storedKubeCred  map[string][]byte
	removedKubeCred []string
}

func (f *fakeJumpPodForJumper) JumpContext() *types.JumpContext {
	return f.jctx
}

func (f *fakeJumpPodForJumper) StoreKubernetesCred(_ context.Context, credID jpent.CredentialID, value []byte) error {
	if f.storedKubeCred == nil {
		f.storedKubeCred = map[string][]byte{}
	}
	f.storedKubeCred[string(credID)] = append([]byte(nil), value...)
	return nil
}

func (f *fakeJumpPodForJumper) RemoveKubernetesCred(_ context.Context, credID jpent.CredentialID) error {
	f.removedKubeCred = append(f.removedKubeCred, string(credID))
	return nil
}

func (f *fakeJumpPodForJumper) TailscaleUp(_ context.Context, _ jmpapi.TailscaleUpRequest) error {
	return nil
}

func (f *fakeJumpPodForJumper) TailscaleDown(_ context.Context) error {
	return nil
}

func (f *fakeJumpPodForJumper) TailscaleStatus(_ context.Context) (*jmpapi.TailscaleStatus, error) {
	return nil, nil
}

func newTestK8sTokenJumperWithConfig(jumpName string) (*K8sTokenJumper, *types.JumpContext, *fake.Clientset) {
	scheme := runtime.NewScheme()
	err := clientgoscheme.AddToScheme(scheme)
	gomega.Expect(err).ToNot(gomega.HaveOccurred())

	ctrlCli := ctrlfake.NewClientBuilder().WithScheme(scheme).Build()
	kubeCli := fake.NewSimpleClientset()

	j := &K8sTokenJumper{
		cfg: clientcmdapi.Config{
			CurrentContext: "base",
			Contexts: map[string]*clientcmdapi.Context{
				"base": {
					Cluster:  "base-cluster",
					AuthInfo: "base-auth",
				},
			},
			Clusters: map[string]*clientcmdapi.Cluster{
				"base-cluster": {
					Server: "https://cluster.example.com",
				},
			},
		},
		cli: &k8sutil.UnionClient{
			Ctrl: ctrlCli,
			Kube: kubeCli,
		},
		roleCreator: fakeJumpRoleCreator{},
	}

	id := types.NewID(jumpName)
	jctx := &types.JumpContext{
		Name:   jumpName,
		ID:     id,
		Namer:  types.NewResourceNamer(id, "jump-system"),
		Logger: logr.Discard(),
	}
	return j, jctx, kubeCli
}

func seedServiceAccountToken(kubeCli *fake.Clientset, namespace, saName, secretName, token string) {
	ctx := context.Background()
	_, err := kubeCli.CoreV1().Secrets(namespace).Create(ctx, &corev1.Secret{
		ObjectMeta: metav1.ObjectMeta{
			Name:      secretName,
			Namespace: namespace,
		},
		Data: map[string][]byte{
			"token": []byte(token),
		},
	}, metav1.CreateOptions{})
	gomega.Expect(err).ToNot(gomega.HaveOccurred())

	_, err = kubeCli.CoreV1().ServiceAccounts(namespace).Create(ctx, &corev1.ServiceAccount{
		ObjectMeta: metav1.ObjectMeta{
			Name:      saName,
			Namespace: namespace,
		},
		Secrets: []corev1.ObjectReference{
			{Name: secretName},
		},
	}, metav1.CreateOptions{})
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
}

type fakeJumpRoleCreator struct{}

func (fakeJumpRoleCreator) PrepareRoles(_ context.Context, _ ctrlclient.Client) error {
	return nil
}

func (fakeJumpRoleCreator) JumpRoleName(priv jpent.JumpPrivilege) string {
	switch priv {
	case jpent.JumpPrivilegeReadOnly:
		return JumpReadOnlyRoleName
	case jpent.JumpPrivilegeReadWrite:
		return JumpReadWriteRoleName
	case jpent.JumpPrivilegeAdmin:
		return JumpAdminRoleName
	default:
		return ""
	}
}

func (fakeJumpRoleCreator) JumpGatewayRoleName(priv jpent.JumpPrivilege) string {
	switch priv {
	case jpent.JumpPrivilegeReadOnly:
		return JumpGatewayReadOnlyRoleName
	case jpent.JumpPrivilegeReadWrite:
		return JumpGatewayReadWriteRoleName
	case jpent.JumpPrivilegeAdmin:
		return JumpGatewayAdminRoleName
	default:
		return ""
	}
}