// Copyright (c) The OpenTofu Authors
// SPDX-License-Identifier: MPL-2.0
// Copyright (c) 2023 HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0

package auth

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"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/azidentity"
	"github.com/opentofu/opentofu/internal/httpclient"
	"github.com/opentofu/opentofu/internal/tfdiags"
)

type OIDCAuthConfig struct {
	UseOIDC           bool
	OIDCToken         string
	OIDCTokenFilePath string
	OIDCRequestURL    string
	OIDCRequestToken  string
}

type oidcAuth struct{}

var _ AuthMethod = &oidcAuth{}

func (cred *oidcAuth) Name() string {
	return "OpenID Connect Auth"
}

func (cred *oidcAuth) Construct(ctx context.Context, config *Config) (azcore.TokenCredential, error) {
	clientId, err := consolidateClientId(config)
	if err != nil {
		return nil, err
	}

	return azidentity.NewClientAssertionCredential(
		config.TenantID,
		clientId,
		// The azure sdk calls this callback whenever it needs client assertion
		//
		// Previously, the OIDC token was fetched once and returned statically,
		// which caused failures when using short-lived tokens in azure devops
		// pipelines during long-running operations.
		//
		// By resolving the token dynamically here, we allow the sdk to obtain
		// a fresh OIDC token as needed, enabling proper token refresh behavior.
		func(ctx context.Context) (string, error) {
			client := httpclient.New(ctx)

			if config.OIDCToken == "" && config.OIDCTokenFilePath == "" {
				return getTokenFromRemote(client, config.OIDCAuthConfig, config.CloudConfig)
			}
			return consolidateToken(config)
		},
		&azidentity.ClientAssertionCredentialOptions{
			ClientOptions: clientOptions(httpclient.New(ctx), config.CloudConfig),
		},
	)
}

type TokenResponse struct {
	Value string `json:"value"`
}

func getTokenFromRemote(client *http.Client, config OIDCAuthConfig, env cloud.Configuration) (string, error) {
	// GET from the request URL, using the bearer token
	req, err := http.NewRequest(http.MethodGet, config.OIDCRequestURL, nil)
	if err != nil {
		return "", fmt.Errorf("malformed token request: %w", err)
	}
	req.Header.Add("Authorization", "Bearer "+config.OIDCRequestToken)
	req.Header.Add("Accept", "application/json; api-version=2.0")
	req.Header.Add("Content-Type", "application/json")

	if !req.URL.Query().Has("audience") {
		query := req.URL.Query()
		query.Set("audience", requestURLAudience(env))
		req.URL.RawQuery = query.Encode()
	}

	// Read the response
	resp, err := client.Do(req)
	if err != nil {
		return "", fmt.Errorf("error obtaining token: %w", err)
	}
	defer resp.Body.Close()
	rawToken, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("io error reading token response body: %w", err)
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return "", fmt.Errorf("non-2xx response: status code %d, body: %s", resp.StatusCode, rawToken)
	}
	var token TokenResponse
	// Provide that response as the access token.
	err = json.Unmarshal(rawToken, &token)
	if err != nil {
		return "", fmt.Errorf("error parsing json of token response body: %w", err)
	}
	return token.Value, nil
}

// requestURLAudience ensures the audience is configured appropriately for
// its target cloud authority
//
// Reference:
// https://learn.microsoft.com/en-us/entra/workload-id/workload-identity-federation-config-app-trust-managed-identity
func requestURLAudience(env cloud.Configuration) string {
	switch env.ActiveDirectoryAuthorityHost {
	case "https://login.chinacloudapi.cn/":
		return "api://AzureADTokenExchangeChina"
	case "https://login.microsoftonline.us/":
		return "api://AzureADTokenExchangeUSGov"
	default:
		return "api://AzureADTokenExchange"
	}
}

func consolidateToken(config *Config) (string, error) {
	return consolidateFileAndValue(config.OIDCToken, config.OIDCTokenFilePath, "token", true)
}

func (cred *oidcAuth) Validate(ctx context.Context, config *Config) tfdiags.Diagnostics {
	var diags tfdiags.Diagnostics
	if !config.UseOIDC {
		diags = diags.Append(tfdiags.Sourceless(
			tfdiags.Error,
			"Invalid Azure OpenID Connect Auth",
			"OpenID Connect Auth is disabled when use_oidc or the environment variable ARM_USE_OIDC are unset or set explicitly to false.",
		))
		return diags
	}
	if config.TenantID == "" {
		diags = diags.Append(tfdiags.Sourceless(
			tfdiags.Error,
			"Invalid Azure OpenID Connect Auth",
			"Tenant ID is missing.",
		))
	}
	_, err := consolidateClientId(config)
	if err != nil {
		diags = diags.Append(tfdiags.Sourceless(
			tfdiags.Error,
			"Invalid Azure OpenID Connect Auth",
			fmt.Sprintf("The Client ID is misconfigured: %s.", tfdiags.FormatError(err)),
		))
	}
	directTokenUnset := config.OIDCToken == "" && config.OIDCTokenFilePath == ""
	indirectTokenUnset := config.OIDCRequestURL == "" || config.OIDCRequestToken == ""
	if directTokenUnset && indirectTokenUnset {
		diags = diags.Append(tfdiags.Sourceless(
			tfdiags.Error,
			"Invalid Azure OpenID Connect Auth",
			"An access token must be provided, either directly with a variable or through a file, or indirectly through a request URL and request token (as in GitHub Actions).",
		))
	}
	if directTokenUnset {
		// check request URL and token
		_, err := getTokenFromRemote(httpclient.New(ctx), config.OIDCAuthConfig, config.CloudConfig)
		if err != nil {
			diags = diags.Append(tfdiags.Sourceless(
				tfdiags.Error,
				"Invalid Azure OpenID Connect Auth",
				fmt.Sprintf("Tried to test fetching the token, but received this error instead: %s.", tfdiags.FormatError(err)),
			))
		}
	}
	// This will work, even if both token and file path are empty
	if _, err := consolidateToken(config); err != nil {
		diags = diags.Append(tfdiags.Sourceless(
			tfdiags.Error,
			"Invalid Azure OpenID Connect Auth",
			fmt.Sprintf("The token is misconfigured: %s", err.Error()),
		))
	}
	return diags
}

func (cred *oidcAuth) AugmentConfig(_ context.Context, config *Config) error {
	return checkNamesForAccessKeyCredentials(config.StorageAddresses)
}
