package k8scliv1

import (
	"context"
	"fmt"

	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v6"
	aksv1beta1 "github.com/upbound/provider-azure/apis/containerservice/v1beta1"
	azurev1beta1 "github.com/upbound/provider-azure/apis/v1beta1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	workloadIdentityLoginMethod = "workloadidentity"
	getTokenCommand             = "get-token"
	argLoginMethod              = "--login"
	argServerID                 = "--server-id"
)

type aksTenantClientManager struct {
	cli    client.Client
	option client.Options
}

func NewAKSClientManager(cli client.Client, opt client.Options) ClientManager {
	return &aksTenantClientManager{
		cli:    cli,
		option: opt,
	}
}

func (m *aksTenantClientManager) GetClient(ctx context.Context, cluster client.ObjectKey) (client.Client, error) {
	cfg, err := m.GetConfig(ctx, cluster)
	if err != nil {
		return nil, err
	}
	return client.New(cfg, m.option)
}

func (m *aksTenantClientManager) GetConfig(ctx context.Context, clusterKey client.ObjectKey) (*rest.Config, error) {
	aks := &aksv1beta1.KubernetesCluster{}
	if err := m.cli.Get(ctx, clusterKey, aks); err != nil {
		return nil, fmt.Errorf("failed to get KubernetesCluster %s: %w", clusterKey, err)
	}

	if aks.Spec.ForProvider.ResourceGroupName == nil {
		return nil, fmt.Errorf("KubernetesCluster %s has no resourceGroupName", clusterKey)
	}
	resourceGroup := *aks.Spec.ForProvider.ResourceGroupName

	var subscriptionID string
	if aks.Spec.ProviderConfigReference != nil && aks.Spec.ProviderConfigReference.Name != "" {
		pc := &azurev1beta1.ProviderConfig{}
		if err := m.cli.Get(ctx, client.ObjectKey{Name: aks.Spec.ProviderConfigReference.Name}, pc); err != nil {
			return nil, fmt.Errorf("failed to get ProviderConfig %s: %w", aks.Spec.ProviderConfigReference.Name, err)
		}
		if pc.Spec.SubscriptionID != nil {
			subscriptionID = *pc.Spec.SubscriptionID
		}
	}
	if subscriptionID == "" {
		return nil, fmt.Errorf("no subscriptionID found for KubernetesCluster %s", clusterKey)
	}

	rawKubeconfig, err := fetchAKSKubeconfig(ctx, subscriptionID, resourceGroup, aks.Name)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch AKS kubeconfig for %s: %w", clusterKey, err)
	}

	wiKubeconfig, err := convertToWorkloadIdentityLogin(rawKubeconfig)
	if err != nil {
		return nil, fmt.Errorf("failed to convert AKS kubeconfig to workload identity: %w", err)
	}

	return clientcmd.RESTConfigFromKubeConfig(wiKubeconfig)
}

// fetchAKSKubeconfig retrieves the raw user kubeconfig for an AKS cluster via
// the Azure ARM API. It relies on workload identity credentials being available
// in the environment (set by the AKS workload identity webhook).
func fetchAKSKubeconfig(ctx context.Context, subscriptionID, resourceGroup, clusterName string) ([]byte, error) {
	cred, err := azidentity.NewDefaultAzureCredential(nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create Azure credential: %w", err)
	}

	armClient, err := armcontainerservice.NewManagedClustersClient(subscriptionID, cred, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create ARM managed clusters client: %w", err)
	}

	resp, err := armClient.ListClusterUserCredentials(ctx, resourceGroup, clusterName, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to list cluster user credentials: %w", err)
	}

	if len(resp.Kubeconfigs) == 0 {
		return nil, fmt.Errorf("no kubeconfig returned for AKS cluster %s", clusterName)
	}

	return resp.Kubeconfigs[0].Value, nil
}

// convertToWorkloadIdentityLogin takes a raw AKS kubeconfig and rewrites the
// exec args to use kubelogin workload-identity authentication.
func convertToWorkloadIdentityLogin(data []byte) ([]byte, error) {
	kubeconfig, err := clientcmd.Load(data)
	if err != nil {
		return nil, fmt.Errorf("failed to load kubeconfig: %w", err)
	}

	if len(kubeconfig.AuthInfos) == 0 {
		return nil, fmt.Errorf("no authInfos found in kubeconfig")
	}

	for _, authInfo := range kubeconfig.AuthInfos {
		original := authInfo.DeepCopy()
		exec := authInfo.Exec
		exec.Args = []string{getTokenCommand}
		exec.Args = append(exec.Args, argLoginMethod, workloadIdentityLoginMethod)

		// all login methods require --server-id specified
		argServerIDVal := getExecArg(original, argServerID)
		if argServerIDVal == "" {
			return nil, fmt.Errorf("%s is required", argServerID)
		}
		exec.Args = append(exec.Args, argServerID, argServerIDVal)
		authInfo.Exec = exec
	}

	modifiedKubeconfig, err := clientcmd.Write(*kubeconfig)
	if err != nil {
		return nil, fmt.Errorf("failed to write kubeconfig: %w", err)
	}
	return modifiedKubeconfig, nil
}

// getExecArg returns the value after someArg in authInfo.Exec.Args.
func getExecArg(authInfo *clientcmdapi.AuthInfo, someArg string) string {
	if someArg == "" || authInfo == nil || authInfo.Exec == nil {
		return ""
	}
	for i, arg := range authInfo.Exec.Args {
		if arg == someArg && i+1 < len(authInfo.Exec.Args) {
			return authInfo.Exec.Args[i+1]
		}
	}
	return ""
}