// Copyright The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package azuread

import (
	"context"
	"errors"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore/cloud"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
	"github.com/google/uuid"
	"github.com/grafana/regexp"
	config_util "github.com/prometheus/common/config"
)

// Clouds.
const (
	AzureChina      = "AzureChina"
	AzureGovernment = "AzureGovernment"
	AzurePublic     = "AzurePublic"
)

// Audiences.
const (
	IngestionChinaAudience      = "https://monitor.azure.cn//.default"
	IngestionGovernmentAudience = "https://monitor.azure.us//.default"
	IngestionPublicAudience     = "https://monitor.azure.com//.default"
)

const (
	// DefaultWorkloadIdentityTokenPath is the default path where the Azure Workload Identity
	// webhook puts the service account token on Azure environments. See <azure docs link>.
	DefaultWorkloadIdentityTokenPath = "/var/run/secrets/azure/tokens/azure-identity-token"
)

// ManagedIdentityConfig is used to store managed identity config values.
type ManagedIdentityConfig struct {
	// ClientID is the clientId of the managed identity that is being used to authenticate.
	ClientID string `yaml:"client_id,omitempty"`
}

// WorkloadIdentityConfig is used to store workload identity config values.
type WorkloadIdentityConfig struct {
	// ClientID is the clientId of the Microsoft Entra application or user-assigned managed identity.
	ClientID string `yaml:"client_id,omitempty"`

	// TenantID is the tenantId of the Microsoft Entra application or user-assigned managed identity.
	// This should match the tenant ID where your application or managed identity is registered.
	TenantID string `yaml:"tenant_id,omitempty"`

	// TokenFilePath is the path to the token file provided by the Kubernetes service account projected volume.
	// If not specified, it defaults to DefaultWorkloadIdentityTokenPath.
	TokenFilePath string `yaml:"token_file_path,omitempty"`
}

// OAuthConfig is used to store azure oauth config values.
type OAuthConfig struct {
	// ClientID is the clientId of the azure active directory application that is being used to authenticate.
	ClientID string `yaml:"client_id,omitempty"`

	// ClientSecret is the clientSecret of the azure active directory application that is being used to authenticate.
	ClientSecret config_util.Secret `yaml:"client_secret,omitempty"`

	// TenantID is the tenantId of the azure active directory application that is being used to authenticate.
	TenantID string `yaml:"tenant_id,omitempty"`
}

// SDKConfig is used to store azure SDK config values.
type SDKConfig struct {
	// TenantID is the tenantId of the azure active directory application that is being used to authenticate.
	TenantID string `yaml:"tenant_id,omitempty"`
}

// AzureADConfig is used to store the config values.
type AzureADConfig struct { //nolint:revive // exported.
	// ManagedIdentity is the managed identity that is being used to authenticate.
	ManagedIdentity *ManagedIdentityConfig `yaml:"managed_identity,omitempty"`

	// WorkloadIdentity is the workload identity that is being used to authenticate.
	WorkloadIdentity *WorkloadIdentityConfig `yaml:"workload_identity,omitempty"`

	// OAuth is the oauth config that is being used to authenticate.
	OAuth *OAuthConfig `yaml:"oauth,omitempty"`

	// SDK is the SDK config that is being used to authenticate.
	SDK *SDKConfig `yaml:"sdk,omitempty"`

	// Cloud is the Azure cloud in which the service is running. Example: AzurePublic/AzureGovernment/AzureChina.
	Cloud string `yaml:"cloud,omitempty"`

	// Scope is the custom OAuth 2.0 scope to request when acquiring tokens.
	Scope string `yaml:"scope,omitempty"`
}

// azureADRoundTripper is used to store the roundtripper and the tokenprovider.
type azureADRoundTripper struct {
	next          http.RoundTripper
	tokenProvider *tokenProvider
}

// tokenProvider is used to store and retrieve Azure AD accessToken.
type tokenProvider struct {
	// token is member used to store the current valid accessToken.
	token string
	// mu guards access to token.
	mu sync.Mutex
	// refreshTime is used to store the refresh time of the current valid accessToken.
	refreshTime time.Time
	// credentialClient is the Azure AD credential client that is being used to retrieve accessToken.
	credentialClient azcore.TokenCredential
	options          *policy.TokenRequestOptions
}

// Validate validates config values provided.
func (c *AzureADConfig) Validate() error {
	if c.Cloud == "" {
		c.Cloud = AzurePublic
	}

	if c.Cloud != AzureChina && c.Cloud != AzureGovernment && c.Cloud != AzurePublic {
		return errors.New("must provide a cloud in the Azure AD config")
	}

	authenticators := 0
	if c.ManagedIdentity != nil {
		authenticators++
	}
	if c.WorkloadIdentity != nil {
		authenticators++
	}
	if c.OAuth != nil {
		authenticators++
	}
	if c.SDK != nil {
		authenticators++
	}

	if authenticators == 0 {
		return errors.New("must provide an Azure Managed Identity, Azure Workload Identity, Azure OAuth or Azure SDK in the Azure AD config")
	}
	if authenticators > 1 {
		return errors.New("cannot provide multiple authentication methods in the Azure AD config")
	}

	if c.ManagedIdentity != nil {
		if c.ManagedIdentity.ClientID != "" {
			_, err := uuid.Parse(c.ManagedIdentity.ClientID)
			if err != nil {
				return errors.New("the provided Azure Managed Identity client_id is invalid")
			}
		}
	}

	if c.WorkloadIdentity != nil {
		if c.WorkloadIdentity.ClientID == "" {
			return errors.New("must provide an Azure Workload Identity client_id in the Azure AD config")
		}
		if c.WorkloadIdentity.TenantID == "" {
			return errors.New("must provide an Azure Workload Identity tenant_id in the Azure AD config")
		}

		if _, err := uuid.Parse(c.WorkloadIdentity.ClientID); err != nil {
			return errors.New("the provided Azure Workload Identity client_id is invalid")
		}
		if _, err := uuid.Parse(c.WorkloadIdentity.TenantID); err != nil {
			return errors.New("the provided Azure Workload Identity tenant_id is invalid")
		}

		if c.WorkloadIdentity.TokenFilePath == "" {
			c.WorkloadIdentity.TokenFilePath = DefaultWorkloadIdentityTokenPath
		}
	}

	if c.OAuth != nil {
		if c.OAuth.ClientID == "" {
			return errors.New("must provide an Azure OAuth client_id in the Azure AD config")
		}
		if c.OAuth.ClientSecret == "" {
			return errors.New("must provide an Azure OAuth client_secret in the Azure AD config")
		}
		if c.OAuth.TenantID == "" {
			return errors.New("must provide an Azure OAuth tenant_id in the Azure AD config")
		}

		if _, err := uuid.Parse(c.OAuth.ClientID); err != nil {
			return errors.New("the provided Azure OAuth client_id is invalid")
		}
		if _, err := regexp.MatchString("^[0-9a-zA-Z-.]+$", c.OAuth.TenantID); err != nil {
			return errors.New("the provided Azure OAuth tenant_id is invalid")
		}
	}

	if c.SDK != nil {
		if c.SDK.TenantID != "" {
			if _, err := regexp.MatchString("^[0-9a-zA-Z-.]+$", c.SDK.TenantID); err != nil {
				return errors.New("the provided Azure SDK tenant_id is invalid")
			}
		}
	}

	if c.Scope != "" {
		if matched, err := regexp.MatchString("^[\\w\\s:/.\\-]+$", c.Scope); err != nil || !matched {
			return errors.New("the provided scope contains invalid characters")
		}
	}

	return nil
}

// UnmarshalYAML unmarshal the Azure AD config yaml.
func (c *AzureADConfig) UnmarshalYAML(unmarshal func(any) error) error {
	type plain AzureADConfig
	*c = AzureADConfig{}
	if err := unmarshal((*plain)(c)); err != nil {
		return err
	}
	return c.Validate()
}

// NewAzureADRoundTripper creates round tripper adding Azure AD authorization to calls.
func NewAzureADRoundTripper(cfg *AzureADConfig, next http.RoundTripper) (http.RoundTripper, error) {
	if next == nil {
		next = http.DefaultTransport
	}

	cred, err := newTokenCredential(cfg)
	if err != nil {
		return nil, err
	}

	tokenProvider, err := newTokenProvider(cfg, cred)
	if err != nil {
		return nil, err
	}

	rt := &azureADRoundTripper{
		next:          next,
		tokenProvider: tokenProvider,
	}
	return rt, nil
}

// RoundTrip sets Authorization header for requests.
func (rt *azureADRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	accessToken, err := rt.tokenProvider.getAccessToken(req.Context())
	if err != nil {
		return nil, err
	}
	bearerAccessToken := "Bearer " + accessToken
	req.Header.Set("Authorization", bearerAccessToken)

	return rt.next.RoundTrip(req)
}

// newTokenCredential returns a TokenCredential of different kinds like Azure Managed Identity, Workload Identity and Azure AD application.
func newTokenCredential(cfg *AzureADConfig) (azcore.TokenCredential, error) {
	var cred azcore.TokenCredential
	var err error
	cloudConfiguration, err := getCloudConfiguration(cfg.Cloud)
	if err != nil {
		return nil, err
	}
	clientOpts := &azcore.ClientOptions{
		Cloud: cloudConfiguration,
	}

	if cfg.ManagedIdentity != nil {
		managedIdentityConfig := &ManagedIdentityConfig{
			ClientID: cfg.ManagedIdentity.ClientID,
		}
		cred, err = newManagedIdentityTokenCredential(clientOpts, managedIdentityConfig)
		if err != nil {
			return nil, err
		}
	}

	if cfg.WorkloadIdentity != nil {
		workloadIdentityConfig := &WorkloadIdentityConfig{
			ClientID:      cfg.WorkloadIdentity.ClientID,
			TenantID:      cfg.WorkloadIdentity.TenantID,
			TokenFilePath: cfg.WorkloadIdentity.TokenFilePath,
		}
		cred, err = newWorkloadIdentityTokenCredential(clientOpts, workloadIdentityConfig)
		if err != nil {
			return nil, err
		}
	}

	if cfg.OAuth != nil {
		oAuthConfig := &OAuthConfig{
			ClientID:     cfg.OAuth.ClientID,
			ClientSecret: cfg.OAuth.ClientSecret,
			TenantID:     cfg.OAuth.TenantID,
		}
		cred, err = newOAuthTokenCredential(clientOpts, oAuthConfig)
		if err != nil {
			return nil, err
		}
	}

	if cfg.SDK != nil {
		sdkConfig := &SDKConfig{
			TenantID: cfg.SDK.TenantID,
		}
		cred, err = newSDKTokenCredential(clientOpts, sdkConfig)
		if err != nil {
			return nil, err
		}
	}

	return cred, nil
}

// newManagedIdentityTokenCredential returns new Managed Identity token credential.
func newManagedIdentityTokenCredential(clientOpts *azcore.ClientOptions, managedIdentityConfig *ManagedIdentityConfig) (azcore.TokenCredential, error) {
	var opts *azidentity.ManagedIdentityCredentialOptions
	if managedIdentityConfig.ClientID != "" {
		clientID := azidentity.ClientID(managedIdentityConfig.ClientID)
		opts = &azidentity.ManagedIdentityCredentialOptions{ClientOptions: *clientOpts, ID: clientID}
	} else {
		opts = &azidentity.ManagedIdentityCredentialOptions{ClientOptions: *clientOpts}
	}
	return azidentity.NewManagedIdentityCredential(opts)
}

// newWorkloadIdentityTokenCredential returns new Microsoft Entra Workload Identity token credential.
//
// For detailed setup instructions, see:
// https://learn.microsoft.com/en-us/azure/azure-monitor/essentials/prometheus-metrics-enable-workload-identity
func newWorkloadIdentityTokenCredential(clientOpts *azcore.ClientOptions, workloadIdentityConfig *WorkloadIdentityConfig) (azcore.TokenCredential, error) {
	opts := &azidentity.WorkloadIdentityCredentialOptions{
		ClientOptions: *clientOpts,
		ClientID:      workloadIdentityConfig.ClientID,
		TenantID:      workloadIdentityConfig.TenantID,
		TokenFilePath: workloadIdentityConfig.TokenFilePath,
	}

	return azidentity.NewWorkloadIdentityCredential(opts)
}

// newOAuthTokenCredential returns new OAuth token credential.
func newOAuthTokenCredential(clientOpts *azcore.ClientOptions, oAuthConfig *OAuthConfig) (azcore.TokenCredential, error) {
	opts := &azidentity.ClientSecretCredentialOptions{ClientOptions: *clientOpts}
	return azidentity.NewClientSecretCredential(oAuthConfig.TenantID, oAuthConfig.ClientID, string(oAuthConfig.ClientSecret), opts)
}

// newSDKTokenCredential returns new SDK token credential.
func newSDKTokenCredential(clientOpts *azcore.ClientOptions, sdkConfig *SDKConfig) (azcore.TokenCredential, error) {
	opts := &azidentity.DefaultAzureCredentialOptions{ClientOptions: *clientOpts, TenantID: sdkConfig.TenantID}
	return azidentity.NewDefaultAzureCredential(opts)
}

// newTokenProvider helps to fetch accessToken for different types of credential. This also takes care of
// refreshing the accessToken before expiry. This accessToken is attached to the Authorization header while making requests.
func newTokenProvider(cfg *AzureADConfig, cred azcore.TokenCredential) (*tokenProvider, error) {
	var scopes []string

	// Use custom scope if provided, otherwise fallback to cloud-specific audience
	if cfg.Scope != "" {
		scopes = []string{cfg.Scope}
	} else {
		audience, err := getAudience(cfg.Cloud)
		if err != nil {
			return nil, err
		}
		scopes = []string{audience}
	}

	tokenProvider := &tokenProvider{
		credentialClient: cred,
		options:          &policy.TokenRequestOptions{Scopes: scopes},
	}

	return tokenProvider, nil
}

// getAccessToken returns the current valid accessToken.
func (tokenProvider *tokenProvider) getAccessToken(ctx context.Context) (string, error) {
	tokenProvider.mu.Lock()
	defer tokenProvider.mu.Unlock()
	if tokenProvider.valid() {
		return tokenProvider.token, nil
	}
	err := tokenProvider.getToken(ctx)
	if err != nil {
		return "", errors.New("Failed to get access token: " + err.Error())
	}
	return tokenProvider.token, nil
}

// valid checks if the token in the token provider is valid and not expired.
func (tokenProvider *tokenProvider) valid() bool {
	if tokenProvider.token == "" {
		return false
	}
	if tokenProvider.refreshTime.After(time.Now().UTC()) {
		return true
	}
	return false
}

// getToken retrieves a new accessToken and stores the newly retrieved token in the tokenProvider.
func (tokenProvider *tokenProvider) getToken(ctx context.Context) error {
	accessToken, err := tokenProvider.credentialClient.GetToken(ctx, *tokenProvider.options)
	if err != nil {
		return err
	}
	if accessToken.Token == "" {
		return errors.New("access token is empty")
	}

	tokenProvider.token = accessToken.Token
	err = tokenProvider.updateRefreshTime(accessToken)
	if err != nil {
		return err
	}
	return nil
}

// updateRefreshTime handles logic to set refreshTime. The refreshTime is set at half the duration of the actual token expiry.
func (tokenProvider *tokenProvider) updateRefreshTime(accessToken azcore.AccessToken) error {
	tokenExpiryTimestamp := accessToken.ExpiresOn.UTC()
	deltaExpirytime := time.Now().Add(time.Until(tokenExpiryTimestamp) / 2)
	if !deltaExpirytime.After(time.Now().UTC()) {
		return errors.New("access token expiry is less than the current time")
	}
	tokenProvider.refreshTime = deltaExpirytime
	return nil
}

// getAudience returns audiences for different clouds.
func getAudience(cloud string) (string, error) {
	switch strings.ToLower(cloud) {
	case strings.ToLower(AzureChina):
		return IngestionChinaAudience, nil
	case strings.ToLower(AzureGovernment):
		return IngestionGovernmentAudience, nil
	case strings.ToLower(AzurePublic):
		return IngestionPublicAudience, nil
	default:
		return "", errors.New("Cloud is not specified or is incorrect: " + cloud)
	}
}

// getCloudConfiguration returns the cloud Configuration which contains AAD endpoint for different clouds.
func getCloudConfiguration(c string) (cloud.Configuration, error) {
	switch strings.ToLower(c) {
	case strings.ToLower(AzureChina):
		return cloud.AzureChina, nil
	case strings.ToLower(AzureGovernment):
		return cloud.AzureGovernment, nil
	case strings.ToLower(AzurePublic):
		return cloud.AzurePublic, nil
	default:
		return cloud.Configuration{}, errors.New("Cloud is not specified or is incorrect: " + c)
	}
}
