package dataplane

import (
	"context"
	"encoding/base64"
	"fmt"
	"os"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/tools/clientcmd"
	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

type EnvProviderManagerOptions struct {
	// Dryrun is used to enable dry run mode which will return a `dryRunClient` k8s client.
	Dryrun bool

	// AssumeRoleName is the name of the AWS IAM role to assume.
	// It's only used for clusters running on AWS.
	AssumeRoleName string

	// ReadonlyAssumeRoleName is the name of the AWS IAM role to assume for readonly access.
	// It's only used for clusters running on AWS.
	// It's required when Dryrun is true and try to get a EnvProvider for clusters running on AWS.
	ReadonlyAssumeRoleName string

	// AssumeRoleExternalID is the external ID of the AWS IAM role to assume.
	// It's only used for clusters running on AWS.
	AssumeRoleExternalID string

	// DataplaneKubeconfig is the path to the kubeconfig file for the dataplane.
	// It should only be used for `update-checker` and local development.
	// Attention: Don't use it if you don't know what you are doing.
	DataplaneKubeconfig string
}

type cachedK8sProvider struct {
	m    K8sClientManager
	cli  client.Client
	opts *EnvProviderManagerOptions
}

func NewEnvProviderManager(scheme *runtime.Scheme, cli client.Client, opts *EnvProviderManagerOptions) *cachedK8sProvider {
	m := NewK8sClientManager(scheme)
	if opts == nil {
		opts = &EnvProviderManagerOptions{}
	}
	return &cachedK8sProvider{m: m, cli: cli, opts: opts}
}

func (p *cachedK8sProvider) GetClient(ctx context.Context, cluster *cloudv1.Cluster) (client.Client, error) {
	provider, err := p.GetEnvProvider(ctx, cluster)
	if err != nil {
		return nil, fmt.Errorf("cannot create dataplane provider: %w", err)
	}

	return provider.GetClient(), nil
}

func (p *cachedK8sProvider) GetEnvProvider(ctx context.Context, cluster *cloudv1.Cluster) (EnvProvider, error) {
	logger := log.FromContext(ctx)
	assumedManageIAMRoleName := p.opts.AssumeRoleName
	if p.opts.Dryrun {
		assumedManageIAMRoleName = p.opts.ReadonlyAssumeRoleName
	}

	// if kubeconfigSecretRef is set, get the kubeconfig from the secret
	var dataplaneKubeData []byte
	switch {
	case p.opts.DataplaneKubeconfig != "":
		data, err := GetDataplaneKubeConfigFromFile(ctx, p.opts.DataplaneKubeconfig)
		if err != nil {
			return nil, fmt.Errorf("cannot get kubeconfig from file %s: %w", p.opts.DataplaneKubeconfig, err)
		}
		logger.Info("get kubeconfig from file", "kubeconfig", p.opts.DataplaneKubeconfig)
		dataplaneKubeData = data
	case cluster.Spec.Kubernetes.KubeConfigSecretRef != nil:
		ref := cluster.Spec.Kubernetes.KubeConfigSecretRef
		data, err := GetDataplaneKubeConfigFromMetaSecret(ctx, p.cli, ref)
		if err != nil {
			return nil, fmt.Errorf("cannot get kubeconfig: %w", err)
		}
		if cluster.GetProviderType() == consts.ProviderTypeAKS {
			// Override kubeconfig with azurecli kubelogin method for local test
			// when environment variable AZURE_FEDERATED_TOKEN_FILE is not set
			data, err = OverrideAksKubeconfig(data)
			if err != nil {
				return nil, fmt.Errorf("cannot override kubeconfig: %w", err)
			}
		}
		dataplaneKubeData = data
	}

	return newEnvProvider(
		p.m,
		cluster,
		assumedManageIAMRoleName,
		p.opts.AssumeRoleExternalID,
		dataplaneKubeData,
		p.opts.Dryrun,
	)
}

func (p *cachedK8sProvider) Deregister(ctx context.Context, cluster *cloudv1.Cluster) {
	p.m.Deregister(cluster.Name)
}

func OverrideAksKubeconfigWithBase64(data []byte) ([]byte, error) {
	// decode data from base64 to string
	data, err := base64.StdEncoding.DecodeString(string(data))
	if err != nil {
		return nil, err
	}

	return OverrideAksKubeconfig(data)
}

// OverrideAksKubeconfig overrides aks kubeconfig when not running with workload identity
func OverrideAksKubeconfig(data []byte) ([]byte, error) {
	if aksWorkloadIdentityEnabled() {
		return data, nil
	}

	cfg, err := clientcmd.Load(data)
	if err != nil {
		return nil, err
	}
	cc := cfg.CurrentContext
	var user string
	for name, ctx := range cfg.Contexts {
		if name == cc {
			user = ctx.AuthInfo
		}
	}
	for name, auth := range cfg.AuthInfos {
		if name == user {
			if auth.Exec != nil {
				overrideExecArgs(auth.Exec)
				return clientcmd.Write(*cfg)
			}
		}
	}
	return data, nil
}

func aksWorkloadIdentityEnabled() bool {
	return os.Getenv("AZURE_FEDERATED_TOKEN_FILE") != ""
}

func overrideExecArgs(execConfig *clientcmdapi.ExecConfig) {
	if execConfig.Command != "kubelogin" {
		return
	}
	args := execConfig.Args
	for idx, arg := range execConfig.Args {
		if arg == "workloadidentity" {
			args[idx] = "azurecli"
		}
	}
	execConfig.Args = args
}