// Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.

package storage

import (
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/pingcap/errors"
)

const (
	// envVolcengineRegion is not part of the official OIDC env var contract
	// (only VOLCENGINE_OIDC_ROLE_TRN / VOLCENGINE_OIDC_TOKEN_FILE are required, see
	// https://github.com/volcengine/volcengine-go-sdk/blob/master/docs/EnvironmentVariables.md).
	// It is injected alongside those two by VKE IRSA and, when present, is used to keep
	// talking to the region-scoped STS host for backward compatibility. It is optional.
	envVolcengineRegion        = "REGION"
	envVolcengineOIDCRoleTRN   = "VOLCENGINE_OIDC_ROLE_TRN"
	envVolcengineOIDCTokenFile = "VOLCENGINE_OIDC_TOKEN_FILE"
	// envVolcengineOIDCSTSEndpoint is the official, documented override for the STS endpoint
	// used by AssumeRoleWithOIDC. When set, it always takes precedence.
	envVolcengineOIDCSTSEndpoint = "VOLCENGINE_OIDC_STS_ENDPOINT"

	volcengineSTSEndpointFormat = "sts.%s.volcengineapi.com"
	// defaultVolcengineSTSEndpoint matches the official SDK's DefaultEndpoint: STS is
	// documented as a global Volcengine service, see
	// https://www.volcengine.com/docs/6257/1494877.
	defaultVolcengineSTSEndpoint     = "sts.volcengineapi.com"
	defaultVolcengineDurationSeconds = 3600
	volcengineOIDCAction             = "AssumeRoleWithOIDC"
	volcengineSTSVersion             = "2018-01-01"
	volcengineOIDCProviderName       = "volcengine_oidc_role_trn"

	volcengineOIDCMaxRetries    = 3
	volcengineOIDCRetryInterval = time.Second
	volcengineOIDCExpiryWindow  = 180
)

type volcengineSessionCredentials struct {
	AccessKeyID         string
	SecretAccessKey     string
	SessionToken        string
	ExpirationTimestamp int64
}

type volcengineOIDCCredentialsProvider struct {
	roleTRN         string
	tokenFilePath   string
	roleSessionName string
	durationSeconds int64
	stsEndpoint     string

	expirationTimestamp int64
	sessionCredentials  *volcengineSessionCredentials
	httpClient          *http.Client
}

type volcengineAPIError struct {
	Code    string `json:"Code"`
	Message string `json:"Message"`
}

type volcengineAssumeRoleWithOIDCResponse struct {
	ResponseMetadata struct {
		RequestID string              `json:"RequestId"`
		Error     *volcengineAPIError `json:"Error"`
	} `json:"ResponseMetadata"`
	Result struct {
		Credentials struct {
			AccessKeyID     string `json:"AccessKeyId"`
			SecretAccessKey string `json:"SecretAccessKey"`
			SessionToken    string `json:"SessionToken"`
			Expiration      string `json:"Expiration"`
			ExpiredTime     string `json:"ExpiredTime"`
		} `json:"Credentials"`
	} `json:"Result"`
}

func autoNewVolcengineOIDCCred() (*credentials.Credentials, error) {
	roleTRN := os.Getenv(envVolcengineOIDCRoleTRN)
	tokenFile := os.Getenv(envVolcengineOIDCTokenFile)
	switch {
	case roleTRN != "" && tokenFile != "":
		return createVolcengineOIDCCred(0)
	case roleTRN != "" || tokenFile != "":
		return nil, errors.Errorf("incomplete volcengine OIDC config: both %s and %s must be set",
			envVolcengineOIDCRoleTRN, envVolcengineOIDCTokenFile)
	default:
		return nil, nil
	}
}

func createVolcengineOIDCCred(duration time.Duration) (*credentials.Credentials, error) {
	provider, err := newVolcengineOIDCCredentialsProvider(duration)
	if err != nil {
		return nil, err
	}
	return credentials.NewCredentials(provider), nil
}

func newVolcengineOIDCCredentialsProvider(duration time.Duration) (*volcengineOIDCCredentialsProvider, error) {
	roleTRN := os.Getenv(envVolcengineOIDCRoleTRN)
	if roleTRN == "" {
		return nil, errors.Errorf("the %s is empty", envVolcengineOIDCRoleTRN)
	}
	tokenFilePath := os.Getenv(envVolcengineOIDCTokenFile)
	if tokenFilePath == "" {
		return nil, errors.Errorf("the %s is empty", envVolcengineOIDCTokenFile)
	}

	durationSeconds := int64(defaultVolcengineDurationSeconds)
	if duration > 0 {
		durationSeconds = int64(duration.Seconds())
	}

	transport, _ := CloneDefaultHttpTransport()
	transport.DialContext = (&net.Dialer{
		Timeout: 5 * time.Second,
	}).DialContext

	return &volcengineOIDCCredentialsProvider{
		roleTRN:         roleTRN,
		tokenFilePath:   tokenFilePath,
		roleSessionName: fmt.Sprintf("credentials-go-%d", time.Now().UnixNano()/1000),
		durationSeconds: durationSeconds,
		stsEndpoint:     resolveVolcengineSTSEndpoint(),
		httpClient: &http.Client{
			Timeout:   15 * time.Second,
			Transport: transport,
		},
	}, nil
}

// resolveVolcengineSTSEndpoint picks the AssumeRoleWithOIDC host, preferring (in order):
//  1. VOLCENGINE_OIDC_STS_ENDPOINT, the official documented override.
//  2. A region-scoped host derived from REGION, kept only for backward compatibility with
//     the VKE IRSA runtime, which happens to inject REGION alongside the OIDC variables.
//  3. The official SDK default, a global endpoint, since STS is a global Volcengine service.
func resolveVolcengineSTSEndpoint() string {
	if endpoint := os.Getenv(envVolcengineOIDCSTSEndpoint); endpoint != "" {
		return endpoint
	}
	if region := os.Getenv(envVolcengineRegion); region != "" {
		return fmt.Sprintf(volcengineSTSEndpointFormat, region)
	}
	return defaultVolcengineSTSEndpoint
}

func (provider *volcengineOIDCCredentialsProvider) needUpdateCredential() bool {
	if provider.expirationTimestamp == 0 {
		return true
	}
	return provider.expirationTimestamp-time.Now().Unix() <= volcengineOIDCExpiryWindow
}

func (provider *volcengineOIDCCredentialsProvider) Retrieve() (credentials.Value, error) {
	if provider.sessionCredentials == nil || provider.needUpdateCredential() {
		sessionCredentials, err := provider.getCredentials()
		if err != nil {
			return credentials.Value{}, err
		}
		provider.sessionCredentials = sessionCredentials
		provider.expirationTimestamp = sessionCredentials.ExpirationTimestamp
	}
	return credentials.Value{
		AccessKeyID:     provider.sessionCredentials.AccessKeyID,
		SecretAccessKey: provider.sessionCredentials.SecretAccessKey,
		SessionToken:    provider.sessionCredentials.SessionToken,
		ProviderName:    volcengineOIDCProviderName,
	}, nil
}

func (provider *volcengineOIDCCredentialsProvider) IsExpired() bool {
	return provider.needUpdateCredential()
}

func (provider *volcengineOIDCCredentialsProvider) getCredentials() (*volcengineSessionCredentials, error) {
	var lastErr error
	for attempt := 0; attempt < volcengineOIDCMaxRetries; attempt++ {
		sessionCredentials, retryable, err := provider.getCredentialsOnce()
		if err == nil {
			return sessionCredentials, nil
		}
		lastErr = err
		if !retryable || attempt == volcengineOIDCMaxRetries-1 {
			return nil, err
		}
		time.Sleep(volcengineOIDCRetryInterval)
	}
	return nil, lastErr
}

func (provider *volcengineOIDCCredentialsProvider) getCredentialsOnce() (*volcengineSessionCredentials, bool, error) {
	token, err := os.ReadFile(provider.tokenFilePath)
	if err != nil {
		return nil, false, errors.Trace(err)
	}

	query := url.Values{}
	query.Set("Action", volcengineOIDCAction)
	query.Set("Version", volcengineSTSVersion)
	requestURL := fmt.Sprintf("https://%s/?%s", provider.stsEndpoint, query.Encode())

	form := url.Values{}
	form.Set("RoleTrn", provider.roleTRN)
	form.Set("RoleSessionName", provider.roleSessionName)
	form.Set("OIDCToken", strings.TrimSpace(string(token)))
	form.Set("DurationSeconds", strconv.FormatInt(provider.durationSeconds, 10))

	req, err := http.NewRequest(http.MethodPost, requestURL, strings.NewReader(form.Encode()))
	if err != nil {
		return nil, false, errors.Trace(err)
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := provider.httpClient.Do(req)
	if err != nil {
		return nil, true, errors.Trace(err)
	}
	defer func() {
		_ = resp.Body.Close()
	}()
	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, true, errors.Trace(err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, isRetryableVolcengineStatus(resp.StatusCode), buildVolcengineHTTPError(resp.StatusCode, respBody)
	}

	var parsedResp volcengineAssumeRoleWithOIDCResponse
	if err := json.Unmarshal(respBody, &parsedResp); err != nil {
		return nil, false, errors.Annotate(err, "failed to unmarshal volcengine sts response")
	}
	if parsedResp.ResponseMetadata.Error != nil {
		return nil, false, errors.Errorf("volcengine sts failed, code=%s message=%s request-id=%s",
			parsedResp.ResponseMetadata.Error.Code,
			parsedResp.ResponseMetadata.Error.Message,
			parsedResp.ResponseMetadata.RequestID)
	}

	creds := parsedResp.Result.Credentials
	if creds.AccessKeyID == "" || creds.SecretAccessKey == "" || creds.SessionToken == "" {
		return nil, false, errors.Errorf("volcengine sts returned empty credentials, request-id=%s", parsedResp.ResponseMetadata.RequestID)
	}

	expirationTimestamp := time.Now().Unix() + provider.durationSeconds
	expirationText := creds.Expiration
	if expirationText == "" {
		expirationText = creds.ExpiredTime
	}
	if expirationText != "" {
		expirationTime, err := time.Parse(time.RFC3339, expirationText)
		if err != nil {
			return nil, false, errors.Annotatef(err, "failed to parse volcengine sts expiration time %s", expirationText)
		}
		expirationTimestamp = expirationTime.Unix()
	}

	return &volcengineSessionCredentials{
		AccessKeyID:         creds.AccessKeyID,
		SecretAccessKey:     creds.SecretAccessKey,
		SessionToken:        creds.SessionToken,
		ExpirationTimestamp: expirationTimestamp,
	}, false, nil
}

func isRetryableVolcengineStatus(statusCode int) bool {
	return statusCode == http.StatusTooManyRequests ||
		statusCode == http.StatusBadGateway ||
		statusCode == http.StatusServiceUnavailable ||
		statusCode == http.StatusGatewayTimeout
}

func buildVolcengineHTTPError(statusCode int, body []byte) error {
	var parsedResp volcengineAssumeRoleWithOIDCResponse
	if err := json.Unmarshal(body, &parsedResp); err == nil && parsedResp.ResponseMetadata.Error != nil {
		return errors.Errorf("get volcengine session token failed, status=%d code=%s message=%s request-id=%s",
			statusCode,
			parsedResp.ResponseMetadata.Error.Code,
			parsedResp.ResponseMetadata.Error.Message,
			parsedResp.ResponseMetadata.RequestID)
	}
	return errors.Errorf("get volcengine session token failed, status=%d body-length=%d", statusCode, len(body))
}