package jumper

import (
	"context"
	"fmt"
	"strings"

	cpeksv1beta1 "github.com/crossplane-contrib/provider-aws/apis/eks/v1beta1"
	cpawsv1beta1 "github.com/crossplane-contrib/provider-aws/apis/v1beta1"
	gkev1beta2 "github.com/crossplane-contrib/provider-gcp/apis/container/v1beta2"
	coresvc "github.com/tidbcloud/core-svc/idl/pbgen/proto/core"
	mgmtv1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"k8s.io/client-go/tools/clientcmd/api"
	"sigs.k8s.io/controller-runtime/pkg/client"

	jpent "github.com/tidbcloud/devops/entity/jump"
	"github.com/tidbcloud/devops/pkg/infra/clients/coresvccli"
	"github.com/tidbcloud/devops/pkg/infra/clients/dedicatedcli"
	"github.com/tidbcloud/devops/pkg/infra/errs"
	"github.com/tidbcloud/devops/pkg/infra/utils/container"
	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"
)

// DedicatedJumper knows how to jump to a dedicated cluster.
type DedicatedJumper struct {
	infraCli         dedicatedcli.InfraAPIClient
	coresvcCli       *coresvccli.CoreSvcCli
	k8sJumperFactory *K8sJumperFactory

	newArchAccessRole string
	aksAuth           *k8sutil.AKSAuth
	configs           *container.Map[string, *api.Config] // TODO: clean cache
}

func NewDedicatedJumper(dep JumperDependency) (Jumper, error) {
	j := &DedicatedJumper{
		infraCli:          dep.Infra,
		coresvcCli:        dep.CoreSvc,
		k8sJumperFactory:  dep.K8sJumperFactory,
		newArchAccessRole: dep.Conf.Dedicated.DedicatedNewArchAccessRole,
		aksAuth: &k8sutil.AKSAuth{
			TenantID:           dep.Conf.Dedicated.DedicatedAKSAuth.TenantID,
			ClientID:           dep.Conf.Dedicated.DedicatedAKSAuth.ClientID,
			FederatedTokenFile: dep.Conf.Dedicated.DedicatedAKSAuth.FederatedTokenFilepath,
		},
		configs: container.NewMap[string, *api.Config](),
	}
	return j, nil
}

func (j *DedicatedJumper) PrepareJump(ctx context.Context, jp JumpPodForJumper, idest jpent.IDestination, param *types.PrepareJumpParam) (*types.JumpParam, error) {
	dest, ok := jpent.AsDestinationDedicatedCluster(idest)
	if !ok {
		return nil, fmt.Errorf("invalid destination type %s", idest.Type())
	}

	network, cluster, err := j.loadNetwork(ctx, dest.ClusterID)
	if err != nil {
		return nil, fmt.Errorf("load network: %w", err)
	}
	if network.Name == "" {
		return nil, fmt.Errorf("infra network is not found")
	}

	var jmpParam *types.JumpParam
	err = j.DoWithCachedConfig(ctx, network, cluster, func(ctx context.Context, config *api.Config) error {
		// proxy to k8s jumper
		k8sJumper, err := j.k8sJumperFactory.NewK8sTokenJumper(*config)
		if err != nil {
			return fmt.Errorf("new k8s jumper: %w", err)
		}

		opt := K8sJumpOption{
			DefaultNamespace: pointer.Pointer(fmt.Sprintf("tidb%d", cluster.Id)),
		}
		jmpParam, err = k8sJumper.PrepareJump(ctx, jp, dest, param, opt)
		return err
	})
	return jmpParam, err
}

func (j *DedicatedJumper) FinalizeJump(ctx context.Context, jp JumpPodForJumper, idest jpent.IDestination, attr jpent.JumpAttributes) error {
	dest, ok := jpent.AsDestinationDedicatedCluster(idest)
	if !ok {
		return fmt.Errorf("invalid destination type %s", idest.Type())
	}

	network, cluster, err := j.loadNetwork(ctx, dest.ClusterID)
	if err != nil {
		if errs.IsNotFound(err) {
			jp.JumpContext().Logger.Info("skip finalizing the jump cause infra network is not found", "err", err.Error())
			return nil
		}

		return fmt.Errorf("load network: %w", err)
	}
	if network.Name == "" {
		jp.JumpContext().Logger.Info("skip finalizing the jump cause infra network is not found")
		return nil
	}

	return j.DoWithCachedConfig(ctx, network, cluster, func(ctx context.Context, config *api.Config) error {
		// proxy to k8s jumper
		k8sJumper, err := j.k8sJumperFactory.NewK8sTokenJumper(*config)
		if err != nil {
			return fmt.Errorf("new k8s jumper: %w", err)
		}

		return k8sJumper.FinalizeJump(ctx, jp, dest)
	})
}

func (j *DedicatedJumper) RefreshJump(ctx context.Context, jp JumpPodForJumper, idest jpent.IDestination, attr jpent.JumpAttributes, param *types.RefreshJumpParam) (*types.JumpParam, error) {
	dest, ok := jpent.AsDestinationDedicatedCluster(idest)
	if !ok {
		return nil, fmt.Errorf("invalid destination type %s", idest.Type())
	}

	network, cluster, err := j.loadNetwork(ctx, dest.ClusterID)
	if err != nil {
		return nil, fmt.Errorf("load network: %w", err)
	}
	if network.Name == "" {
		return nil, fmt.Errorf("infra network is not found")
	}

	var jmpParam *types.JumpParam
	err = j.DoWithCachedConfig(ctx, network, cluster, func(ctx context.Context, config *api.Config) error {
		// proxy to k8s jumper
		k8sJumper, err := j.k8sJumperFactory.NewK8sTokenJumper(*config)
		if err != nil {
			return fmt.Errorf("new k8s jumper: %w", err)
		}

		opt := K8sJumpOption{
			DefaultNamespace: pointer.Pointer(fmt.Sprintf("tidb%d", cluster.Id)),
		}
		jmpParam, err = k8sJumper.RefreshJump(ctx, jp, dest, attr, param, opt)
		return err
	})
	return jmpParam, err
}

func (j *DedicatedJumper) SupportedDestinations() (jpent.DestinationType, []string) {
	typ := jpent.DestinationTypeDedicated
	return typ, nil
}

func (j *DedicatedJumper) loadNetwork(ctx context.Context, clusterID uint64) (client.ObjectKey, *coresvc.Cluster, error) {
	clusterInCS, err := j.coresvcCli.GetCluster(ctx, clusterID, true)
	if err != nil {
		return client.ObjectKey{}, nil, fmt.Errorf("get cluster %d from core svc: %w", clusterID, err)
	}
	network := client.ObjectKey{
		Name:      clusterInCS.DataPlaneInfo.NetworkName,
		Namespace: clusterInCS.DataPlaneInfo.K8SNamespace,
	}
	return network, clusterInCS, nil
}

func (j *DedicatedJumper) DoWithCachedConfig(ctx context.Context, networkKey client.ObjectKey, cluster *coresvc.Cluster, do func(ctx context.Context, cfg *api.Config) error) error {
	key := networkKey.String()

	// load or store config
	_, exist := j.configs.Load(key)
	if !exist {
		config, err := j.buildConfig(ctx, networkKey, cluster)
		if err != nil {
			return err
		}
		j.configs.Store(key, config)
	}

	// do
	config, _ := j.configs.Load(key)
	err := do(ctx, config)
	if err != nil {
		j.configs.Delete(key) // clear invalid config
		return err
	}

	return nil
}

func (j *DedicatedJumper) buildConfig(ctx context.Context, networkKey client.ObjectKey, cluster *coresvc.Cluster) (*api.Config, error) {
	var config *api.Config

	newArch := false
	if cluster.Attributes != nil && cluster.Attributes.UseNewDedicatedArch != nil {
		newArch = *cluster.Attributes.UseNewDedicatedArch
	}
	var err error
	switch cluster.InfraProviderType {
	case coresvccli.InfraProviderTypeAKS:
		config, err = BuildAKSConfig(ctx, j.coresvcCli, cluster, j.aksAuth)
	case coresvccli.InfraProviderTypeEKS:
		if newArch {
			config, err = BuildNewArchEKSConfig(ctx, j.coresvcCli, cluster, j.newArchAccessRole)
		} else {
			config, err = BuildEKSConfig(ctx, j.infraCli.GetClient(), networkKey)
		}
	case coresvccli.InfraProviderTypeGKE:
		if newArch {
			config, err = BuildNewArchGKEConfig(ctx, j.coresvcCli, cluster)
		} else {
			config, err = BuildGKEConfig(ctx, j.infraCli.GetClient(), networkKey)
		}
	case coresvccli.InfraProviderTypeAWS, coresvccli.InfraProviderTypeGCP:
		err = fmt.Errorf("shoot cluster is not supported")
	default:
		return nil, fmt.Errorf("unsopprted cluster type %s", cluster.InfraProviderType)
	}
	if err != nil {
		return nil, fmt.Errorf("build kube config: %w", err)
	}

	return config, nil
}

func BuildNewArchEKSConfig(ctx context.Context, cli *coresvccli.CoreSvcCli, cluster *coresvc.Cluster, role string) (*api.Config, error) {
	compatNetwork, err := cli.Mgmt.Dataplane.GetCompatNetwork(ctx, &mgmtv1.GetCompatNetworkRequest{
		ProjectId:    cluster.ProjectId,
		ProviderName: cluster.InfraProviderType,
		RegionName:   cluster.RegionName,
	})
	if err != nil {
		return nil, fmt.Errorf("get compat network failed: %w", err)
	}

	eksAuth := k8sutil.EKSAuth{
		Region:  cluster.RegionName,
		EKSName: compatNetwork.K8SName,
		RoleARN: role,
	}

	config := k8sutil.NewConfigBuilder().
		AddEKSAuth(compatNetwork.K8SName, eksAuth).
		AddInsecureCluster(compatNetwork.K8SName, compatNetwork.K8SEndpoint).
		AddContext(compatNetwork.K8SName, api.Context{Cluster: compatNetwork.K8SName, AuthInfo: compatNetwork.K8SName}).
		CurrentContext(compatNetwork.K8SName).
		Build()
	return config, nil
}

// BuildEKSConfig retrieves metadata in infra and build a kube config used to accesss the eks cluster.
//
// NOTE: This function is tightly coupled with the infra implementation, so be careful when modifying it.
func BuildEKSConfig(ctx context.Context, infraCli client.Client, networkKey client.ObjectKey) (*api.Config, error) {
	eksList := &cpeksv1beta1.ClusterList{}

	err := infraCli.List(ctx, eksList, client.MatchingLabels{
		"tidbcloud-namespace": networkKey.Namespace,
		"tidbcloud-network":   networkKey.Name,
	})
	if err != nil {
		return nil, fmt.Errorf("list eks cluster failed: %w", err)
	}
	if len(eksList.Items) == 0 {
		return nil, fmt.Errorf("eks cluster for network %s not found", networkKey)
	}

	eks := &eksList.Items[0]
	eksAuth := k8sutil.EKSAuth{
		Region:  pointer.Deref(eks.Spec.ForProvider.Region, ""),
		EKSName: eks.Name,
	}

	if eks.Spec.ProviderConfigReference != nil && eks.Spec.ProviderConfigReference.Name != "" {
		providerCfg := &cpawsv1beta1.ProviderConfig{}
		err := infraCli.Get(ctx, client.ObjectKey{Name: eks.Spec.ProviderConfigReference.Name}, providerCfg)
		if err != nil {
			return nil, fmt.Errorf("get provider config %s failed: %w", eks.Spec.ProviderConfigReference.Name, err)
		}
		if providerCfg.Spec.AssumeRoleARN != nil {
			eksAuth.RoleARN = *providerCfg.Spec.AssumeRoleARN
			if providerCfg.Spec.ExternalID != nil {
				eksAuth.ExternalID = *providerCfg.Spec.ExternalID
			}
		}
	}

	config := k8sutil.NewConfigBuilder().
		AddEKSAuth(eks.Name, eksAuth).
		AddInsecureCluster(eks.Name, eks.Status.AtProvider.Endpoint).
		AddContext(eks.Name, api.Context{Cluster: eks.Name, AuthInfo: eks.Name}).
		CurrentContext(eks.Name).
		Build()
	return config, nil
}

func BuildNewArchGKEConfig(ctx context.Context, cli *coresvccli.CoreSvcCli, cluster *coresvc.Cluster) (*api.Config, error) {
	compatNetwork, err := cli.Mgmt.Dataplane.GetCompatNetwork(ctx, &mgmtv1.GetCompatNetworkRequest{
		ProjectId:    cluster.ProjectId,
		ProviderName: cluster.InfraProviderType,
		RegionName:   cluster.RegionName,
	})

	if err != nil {
		return nil, fmt.Errorf("get compat network failed: %w", err)
	}
	config := k8sutil.NewConfigBuilder().
		AddGKEAuth(compatNetwork.K8SName).
		AddInsecureCluster(compatNetwork.K8SName, compatNetwork.K8SEndpoint).
		AddContext(compatNetwork.K8SName, api.Context{Cluster: compatNetwork.K8SName, AuthInfo: compatNetwork.K8SName}).
		CurrentContext(compatNetwork.K8SName).
		Build()
	return config, nil
}

// BuildGKEConfig retrieves metadata in infra and build a kube config used to accesss the gke cluster.
//
// NOTE: This function is tightly coupled with the infra implementation, so be careful when modifying it.
func BuildGKEConfig(ctx context.Context, infraCli client.Client, networkKey client.ObjectKey) (*api.Config, error) {
	gkeList := &gkev1beta2.ClusterList{}
	err := infraCli.List(ctx, gkeList, client.MatchingLabels{
		"tidbcloud-namespace": networkKey.Namespace,
		"tidbcloud-network":   networkKey.Name,
	})
	if err != nil {
		return nil, fmt.Errorf("list eks cluster failed: %w", err)
	}
	if len(gkeList.Items) == 0 {
		return nil, fmt.Errorf("gke cluster for network %s not found", networkKey)
	}

	gke := &gkeList.Items[0]
	config := k8sutil.NewConfigBuilder().
		AddGKEAuth(gke.Name).
		AddInsecureCluster(gke.Name, "https://"+gke.Status.AtProvider.Endpoint).
		AddContext(gke.Name, api.Context{Cluster: gke.Name, AuthInfo: gke.Name}).
		CurrentContext(gke.Name).
		Build()
	return config, nil
}

func BuildAKSConfig(ctx context.Context, cli *coresvccli.CoreSvcCli, cluster *coresvc.Cluster, aksAuth *k8sutil.AKSAuth) (*api.Config, error) {
	compatNetwork, err := cli.Mgmt.Dataplane.GetCompatNetwork(ctx, &mgmtv1.GetCompatNetworkRequest{
		ProjectId:    cluster.ProjectId,
		ProviderName: cluster.InfraProviderType,
		RegionName:   cluster.RegionName,
	})
	if err != nil {
		return nil, fmt.Errorf("get compat network failed: %w", err)
	}

	if !strings.HasPrefix(compatNetwork.K8SEndpoint, "https://") {
		compatNetwork.K8SEndpoint = "https://" + compatNetwork.K8SEndpoint // TODO: add https prefix now, remove it after mgmt api is fixed.
	}

	config := k8sutil.NewConfigBuilder().
		AddAKSAuth(compatNetwork.K8SName, aksAuth).
		AddInsecureCluster(compatNetwork.K8SName, compatNetwork.K8SEndpoint).
		AddContext(compatNetwork.K8SName, api.Context{Cluster: compatNetwork.K8SName, AuthInfo: compatNetwork.K8SName}).
		CurrentContext(compatNetwork.K8SName).
		Build()
	return config, nil
}