package jumper

import (
	"context"
	"errors"
	"fmt"

	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"
	k8stypes "k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/clientcmd/api"
	"sigs.k8s.io/controller-runtime/pkg/client"
	ctrlutil "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"

	jpent "github.com/tidbcloud/devops/entity/jump"
	k8sutil "github.com/tidbcloud/devops/pkg/infra/utils/k8s"
	"github.com/tidbcloud/devops/pkg/infra/utils/pointer"
	"github.com/tidbcloud/devops/pkg/service/jump/jumppod/types"
)

type K8sJumpOption struct {
	DefaultNamespace *string
}

type K8sJumperFactory struct {
	roleCreator JumpRoleCreator
}

func NewK8sJumperFactory(jumpRoleDir string) (*K8sJumperFactory, error) {
	roleCreator, err := NewJumpRoleCreator(jumpRoleDir)
	if err != nil {
		return nil, fmt.Errorf("failed to create jump role creator: %w", err)
	}

	return &K8sJumperFactory{
		roleCreator: roleCreator,
	}, nil
}

func (f *K8sJumperFactory) NewK8sTokenJumper(cfg api.Config) (*K8sTokenJumper, error) {
	restConfig, err := clientcmd.NewDefaultClientConfig(cfg, nil).ClientConfig()
	if err != nil {
		return nil, fmt.Errorf("create rest config: %w", err)
	}

	cli, err := k8sutil.NewUnionClientFromRest(restConfig, nil)
	if err != nil {
		return nil, fmt.Errorf("create client: %w", err)
	}

	return &K8sTokenJumper{
		cfg:         cfg,
		cli:         cli,
		roleCreator: f.roleCreator,
	}, nil
}

// K8sTokenJumper knows how to jump to a k8s cluster.
// It use a service account to generate a token, and store the token to the jump pod.
type K8sTokenJumper struct {
	cfg         api.Config
	cli         *k8sutil.UnionClient
	roleCreator JumpRoleCreator
}

func (j *K8sTokenJumper) PrepareJump(ctx context.Context, jp JumpPodForJumper, dest jpent.IDestination, param *types.PrepareJumpParam, opt K8sJumpOption) (*types.JumpParam, error) {
	err := j.prepareRBAC(ctx, jp.JumpContext(), dest, param)
	if err != nil {
		return nil, fmt.Errorf("prepare rbac failed: %w", err)
	}

	return j.prepareToken(ctx, jp, dest, param.GenerateTokenParam, param.EnableGatewayAccess, opt)
}

func (j *K8sTokenJumper) FinalizeJump(ctx context.Context, jp JumpPodForJumper, dest jpent.IDestination) error {
	// remove credential
	credID := jpent.FormatCredentialID(dest)
	err := jp.RemoveKubernetesCred(ctx, credID)
	if err != nil {
		return fmt.Errorf("delete cred failed: %w", err)
	}

	gatewayCredID := jpent.FormatGatewayCredentialID(dest)
	err = jp.RemoveKubernetesCred(ctx, gatewayCredID)
	if err != nil {
		return fmt.Errorf("delete gateway cred failed: %w", err)
	}

	// delete rbac resources
	err = j.cleanupRBAC(ctx, jp.JumpContext(), dest)
	if err != nil {
		return fmt.Errorf("cleanup rbac failed: %w", err)
	}

	return nil
}

func (j *K8sTokenJumper) RefreshJump(ctx context.Context, jp JumpPodForJumper, dest jpent.IDestination, attr jpent.JumpAttributes, param *types.RefreshJumpParam, opt K8sJumpOption) (*types.JumpParam, error) {
	enableGatewayAccess := attr.K8s != nil && attr.K8s.EnableGatewayAccess
	return j.prepareToken(ctx, jp, dest, param.GenerateTokenParam, enableGatewayAccess, opt)
}

func (j *K8sTokenJumper) prepareToken(ctx context.Context, jp JumpPodForJumper, dest jpent.IDestination, param types.GenerateTokenParam, enableGatewayAccess bool, opt K8sJumpOption) (*types.JumpParam, error) {
	isGenToken, err := k8sutil.IsAPIGroupVersionResourceSupported(j.cli.Kube.Discovery(), "v1", "serviceaccounts/token")
	if err != nil {
		return nil, fmt.Errorf("probe api failed: %w", err)
	}

	// generate token / kubeconfig for primary access
	token, err := j.retrieveToken(ctx, jp.JumpContext().Namer.JumpServiceAccount(), isGenToken, param)
	if err != nil {
		return nil, fmt.Errorf("generate primary token failed: %w", err)
	}
	credID := jpent.FormatCredentialID(dest)
	if err = j.storeTokenAsKubeconfig(ctx, jp, credID, token.Data, opt); err != nil {
		return nil, fmt.Errorf("store primary cred failed: %w", err)
	}

	// optionally generate token / kubeconfig for gateway access
	if enableGatewayAccess {
		gatewayToken, err := j.retrieveToken(ctx, jp.JumpContext().Namer.JumpGatewayServiceAccount(), isGenToken, param)
		if err != nil {
			return nil, fmt.Errorf("generate gateway token failed: %w", err)
		}
		gatewayCredID := jpent.FormatGatewayCredentialID(dest)
		if err = j.storeTokenAsKubeconfig(ctx, jp, gatewayCredID, gatewayToken.Data, opt); err != nil {
			return nil, fmt.Errorf("store gateway cred failed: %w", err)
		}
	}

	return &types.JumpParam{
		Method: jpent.JumpMethodK8s,
		Attr: jpent.JumpAttributes{
			K8s: &jpent.JumpK8sAttributes{
				Context:             string(credID),
				TokenExpireAt:       token.ExpiredAt.Unix(),
				EnableGatewayAccess: enableGatewayAccess,
			},
		},
	}, nil
}

func (j *K8sTokenJumper) retrieveToken(ctx context.Context, saName k8stypes.NamespacedName, isGenToken bool, param types.GenerateTokenParam) (*types.Token, error) {
	if isGenToken {
		token, err := GenerateToken(ctx, j.cli.Kube, saName, param)
		if err != nil {
			return nil, fmt.Errorf("generate token failed: %w", err)
		}
		return token, nil
	}

	token, err := RetrieveToken(ctx, j.cli.Kube, saName, param)
	if err != nil {
		return nil, fmt.Errorf("retrieve token failed: %w", err)
	}
	return token, nil
}

func (j *K8sTokenJumper) storeTokenAsKubeconfig(ctx context.Context, jp JumpPodForJumper, credID jpent.CredentialID, token string, opt K8sJumpOption) error {
	// build a new kubeconfig
	// * reuse the current cluster
	// * use generated token as auth
	// * cluster name and auth name are prefixed with destination location
	// * use destination location as context name
	curClusterName := j.cfg.Contexts[j.cfg.CurrentContext].Cluster
	curAuthName := j.cfg.Contexts[j.cfg.CurrentContext].AuthInfo
	wrapClusterName := fmt.Sprintf("%s:%s", credID, curClusterName)
	wrapAuthName := fmt.Sprintf("%s:%s", credID, curAuthName)
	kubecfg := k8sutil.NewConfigBuilder().
		AddCluster(wrapClusterName, j.cfg.Clusters[curClusterName]).
		AddTokenAuth(wrapAuthName, token).
		AddContext(string(credID), api.Context{Cluster: wrapClusterName, AuthInfo: wrapAuthName, Namespace: pointer.Deref(opt.DefaultNamespace, "default")}).
		CurrentContext(string(credID)).
		Build()
	data, err := clientcmd.Write(*kubecfg)
	if err != nil {
		return fmt.Errorf("marshal kubeconfig: %w", err)
	}

	if err = jp.StoreKubernetesCred(ctx, credID, data); err != nil {
		return fmt.Errorf("store cred failed: %w", err)
	}
	return nil
}

func (j *K8sTokenJumper) prepareRBAC(ctx context.Context, jctx *types.JumpContext, dest jpent.IDestination, param *types.PrepareJumpParam) error {
	commonLabels := types.NewLabels(nil).SetID(jctx.ID).SetName(jctx.Name)
	commonAnnos := types.NewAnnotations(nil).SetJumpDestination(dest.ID())

	jumpNs := &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: jctx.Namer.Namespace}}
	_, err := j.cli.Kube.CoreV1().Namespaces().Create(ctx, jumpNs, metav1.CreateOptions{})
	if err != nil && !apierrors.IsAlreadyExists(err) {
		return fmt.Errorf("create namespace failed: %w", err)
	}

	jumpSA := &corev1.ServiceAccount{
		ObjectMeta: metav1.ObjectMeta{
			Name:      jctx.Namer.JumpServiceAccount().Name,
			Namespace: jctx.Namer.JumpServiceAccount().Namespace,
		},
	}
	_, err = ctrlutil.CreateOrUpdate(ctx, j.cli.Ctrl, jumpSA, func() error {
		if jumpSA.Labels == nil {
			jumpSA.Labels = map[string]string{}
		}
		commonLabels.InjectTo(jumpSA.Labels)

		if jumpSA.Annotations == nil {
			jumpSA.Annotations = map[string]string{}
		}
		commonAnnos.InjectTo(jumpSA.Annotations)
		return nil
	})
	if err != nil {
		return fmt.Errorf("create or update sa %s failed: %w", jumpSA.Name, err)
	}

	err = j.roleCreator.PrepareRoles(ctx, j.cli.Ctrl)
	if err != nil {
		return fmt.Errorf("prepare shared jump role failed: %w", err)
	}

	jumpRoleBinding := &rbacv1.ClusterRoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Name: jctx.Namer.JumpClusterRoleBinding().Name,
		},
	}
	_, err = ctrlutil.CreateOrUpdate(ctx, j.cli.Ctrl, jumpRoleBinding, func() error {
		if jumpRoleBinding.Labels == nil {
			jumpRoleBinding.Labels = map[string]string{}
		}
		commonLabels.InjectTo(jumpRoleBinding.Labels)

		if jumpRoleBinding.Annotations == nil {
			jumpRoleBinding.Annotations = map[string]string{}
		}
		commonAnnos.InjectTo(jumpRoleBinding.Annotations)

		jumpRoleBinding.Subjects = []rbacv1.Subject{
			{
				Kind:      "ServiceAccount",
				Name:      jumpSA.Name,
				Namespace: jumpSA.Namespace,
			},
		}

		jumpRoleBinding.RoleRef = rbacv1.RoleRef{
			APIGroup: "rbac.authorization.k8s.io",
			Kind:     "ClusterRole",
			Name:     OverwriteJumpRoleName(jumpRoleBinding.RoleRef.Name, JumpRoleName(param.Privilege)), // in case the role binding already exists, we only update the role if the new privilege is higher than the old one
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("create or update role binding %s failed: %w", jumpRoleBinding.Name, err)
	}

	if param.EnableGatewayAccess {
		gatewaySA := &corev1.ServiceAccount{
			ObjectMeta: metav1.ObjectMeta{
				Name:      jctx.Namer.JumpGatewayServiceAccount().Name,
				Namespace: jctx.Namer.JumpGatewayServiceAccount().Namespace,
			},
		}
		_, err := ctrlutil.CreateOrUpdate(ctx, j.cli.Ctrl, gatewaySA, func() error {
			if gatewaySA.Labels == nil {
				gatewaySA.Labels = map[string]string{}
			}
			commonLabels.InjectTo(gatewaySA.Labels)

			if gatewaySA.Annotations == nil {
				gatewaySA.Annotations = map[string]string{}
			}
			commonAnnos.InjectTo(gatewaySA.Annotations)
			return nil
		})
		if err != nil {
			return fmt.Errorf("create or update gateway sa %s failed: %w", gatewaySA.Name, err)
		}

		gatewayRoleBinding := &rbacv1.ClusterRoleBinding{
			ObjectMeta: metav1.ObjectMeta{
				Name: jctx.Namer.JumpGatewayClusterRoleBinding().Name,
			},
		}
		_, err = ctrlutil.CreateOrUpdate(ctx, j.cli.Ctrl, gatewayRoleBinding, func() error {
			if gatewayRoleBinding.Labels == nil {
				gatewayRoleBinding.Labels = map[string]string{}
			}
			commonLabels.InjectTo(gatewayRoleBinding.Labels)

			if gatewayRoleBinding.Annotations == nil {
				gatewayRoleBinding.Annotations = map[string]string{}
			}
			commonAnnos.InjectTo(gatewayRoleBinding.Annotations)

			gatewayRoleBinding.Subjects = []rbacv1.Subject{
				{
					Kind:      "ServiceAccount",
					Name:      gatewaySA.Name,
					Namespace: gatewaySA.Namespace,
				},
			}

			gatewayRoleBinding.RoleRef = rbacv1.RoleRef{
				APIGroup: "rbac.authorization.k8s.io",
				Kind:     "ClusterRole",
				Name:     JumpGatewayRoleName(jpent.JumpPrivilegeReadOnly), // only support readonly now.
			}
			return nil
		})
		if err != nil {
			return fmt.Errorf("create or update gateway role binding %s failed: %w", gatewayRoleBinding.Name, err)
		}
	}

	return nil
}

func (j *K8sTokenJumper) cleanupRBAC(ctx context.Context, jctx *types.JumpContext, dest jpent.IDestination) error {
	jumpSA := &corev1.ServiceAccount{
		ObjectMeta: metav1.ObjectMeta{
			Name:      jctx.Namer.JumpServiceAccount().Name,
			Namespace: jctx.Namer.JumpServiceAccount().Namespace,
		},
	}
	jumpLegacyRole := &rbacv1.ClusterRole{
		ObjectMeta: metav1.ObjectMeta{
			Name:      jctx.Namer.JumpClusterRole().Name,      // nolint:staticcheck
			Namespace: jctx.Namer.JumpClusterRole().Namespace, // nolint:staticcheck
		},
	}
	jumpRoleBinding := &rbacv1.ClusterRoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Name:      jctx.Namer.JumpClusterRoleBinding().Name,
			Namespace: jctx.Namer.JumpClusterRoleBinding().Namespace,
		},
	}
	jumpGatewaySA := &corev1.ServiceAccount{
		ObjectMeta: metav1.ObjectMeta{
			Name:      jctx.Namer.JumpGatewayServiceAccount().Name,
			Namespace: jctx.Namer.JumpGatewayServiceAccount().Namespace,
		},
	}
	jumpGatewayLegacyRole := &rbacv1.ClusterRole{
		ObjectMeta: metav1.ObjectMeta{
			Name:      jctx.Namer.JumpGatewayClusterRole().Name,      // nolint:staticcheck
			Namespace: jctx.Namer.JumpGatewayClusterRole().Namespace, // nolint:staticcheck
		},
	}
	jumpGatewayRoleBinding := &rbacv1.ClusterRoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Name:      jctx.Namer.JumpGatewayClusterRoleBinding().Name,
			Namespace: jctx.Namer.JumpGatewayClusterRoleBinding().Namespace,
		},
	}

	objs := []client.Object{
		jumpSA, jumpLegacyRole, jumpRoleBinding,
		jumpGatewaySA, jumpGatewayLegacyRole, jumpGatewayRoleBinding,
	}
	var errs []error
	for _, obj := range objs {
		err := j.cleanResource(ctx, obj, dest)
		if err != nil {
			errs = append(errs, err)
		}
	}

	return errors.Join(errs...)
}

func (j *K8sTokenJumper) cleanResource(ctx context.Context, obj client.Object, dest jpent.IDestination) error {
	err := j.cli.Ctrl.Get(ctx, client.ObjectKeyFromObject(obj), obj)
	if err != nil && !apierrors.IsNotFound(err) {
		return err
	}

	if apierrors.IsNotFound(err) {
		return nil
	}

	annos := types.NewAnnotations(obj.GetAnnotations())
	if destID, exist := annos.GetJumpDestination(); exist && destID != "" {
		if destID != dest.ID() {
			// not belong to the destination
			return nil
		}
	}

	return j.cli.Ctrl.Delete(ctx, obj)
}