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

package storage

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

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

const (
	envTKERegion               = "TKE_REGION"
	envTKEProviderID           = "TKE_PROVIDER_ID"
	envTKEWebIdentityTokenFile = "TKE_WEB_IDENTITY_TOKEN_FILE"
	envTKERoleARN              = "TKE_ROLE_ARN"

	defaultTencentSTSEndpoint = "sts.tencentcloudapi.com"
	tencentOIDCAction         = "AssumeRoleWithWebIdentity"
	tencentVersion            = "2018-08-13"
	tencentOIDCProviderName   = "tencentcloud_oidc_role_arn"

	defaultTencentDurationSeconds = 3600 * 12
	minTencentDurationSeconds     = 900
)

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

type tencentCloudOIDCCredentialsProvider struct {
	region        string
	providerID    string
	tokenFilePath string
	roleARN       string

	roleSessionName string
	durationSeconds int64
	stsEndpoint     string

	expirationTimestamp int64
	sessionCredentials  *tencentCloudSessionCredentials
	httpClient          *http.Client
}

type tencentCloudAssumeRoleRequest struct {
	ProviderID       string `json:"ProviderId"`
	WebIdentityToken string `json:"WebIdentityToken"`
	RoleArn          string `json:"RoleArn"`
	RoleSessionName  string `json:"RoleSessionName"`
	DurationSeconds  int64  `json:"DurationSeconds"`
}

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

type tencentCloudAssumeRoleResponse struct {
	Response struct {
		Credentials struct {
			Token        string `json:"Token"`
			TmpSecretID  string `json:"TmpSecretId"`
			TmpSecretKey string `json:"TmpSecretKey"`
		} `json:"Credentials"`
		ExpiredTime int64                 `json:"ExpiredTime"`
		Error       *tencentCloudAPIError `json:"Error"`
		RequestID   string                `json:"RequestId"`
	} `json:"Response"`
}

func newTencentCloudOIDCCredentialsProvider(duration time.Duration) (*tencentCloudOIDCCredentialsProvider, error) {
	region := os.Getenv(envTKERegion)
	if len(region) == 0 {
		return nil, errors.Errorf("the %s is empty", envTKERegion)
	}
	providerID := os.Getenv(envTKEProviderID)
	if len(providerID) == 0 {
		return nil, errors.Errorf("the %s is empty", envTKEProviderID)
	}
	tokenFilePath := os.Getenv(envTKEWebIdentityTokenFile)
	if len(tokenFilePath) == 0 {
		return nil, errors.Errorf("the %s is empty", envTKEWebIdentityTokenFile)
	}
	roleARN := os.Getenv(envTKERoleARN)
	if len(roleARN) == 0 {
		return nil, errors.Errorf("the %s is empty", envTKERoleARN)
	}

	durationSeconds := int64(defaultTencentDurationSeconds)
	if duration > 0 {
		durationSeconds = int64(duration.Seconds())
	}
	if durationSeconds < minTencentDurationSeconds {
		return nil, errors.Errorf("the Assume Role session duration of Tencent Cloud should be >= %d seconds", minTencentDurationSeconds)
	}
	transport, _ := CloneDefaultHttpTransport()
	transport.DialContext = (&net.Dialer{
		Timeout: 5 * time.Second,
	}).DialContext

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

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

func (provider *tencentCloudOIDCCredentialsProvider) 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:    tencentOIDCProviderName,
	}, nil
}

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

func (provider *tencentCloudOIDCCredentialsProvider) getCredentials() (*tencentCloudSessionCredentials, error) {
	token, err := os.ReadFile(provider.tokenFilePath)
	if err != nil {
		return nil, errors.Trace(err)
	}

	reqBody := tencentCloudAssumeRoleRequest{
		ProviderID:       provider.providerID,
		WebIdentityToken: strings.TrimSpace(string(token)),
		RoleArn:          provider.roleARN,
		RoleSessionName:  provider.roleSessionName,
		DurationSeconds:  provider.durationSeconds,
	}
	payload, err := json.Marshal(reqBody)
	if err != nil {
		return nil, errors.Trace(err)
	}

	requestURL := fmt.Sprintf("https://%s", provider.stsEndpoint)
	req, err := http.NewRequest(http.MethodPost, requestURL, bytes.NewReader(payload))
	if err != nil {
		return nil, errors.Trace(err)
	}
	req.Header.Set("Authorization", "SKIP")
	req.Header.Set("Content-Type", "application/json; charset=utf-8")
	req.Header.Set("Host", provider.stsEndpoint)
	req.Header.Set("X-TC-Action", tencentOIDCAction)
	req.Header.Set("X-TC-Version", tencentVersion)
	req.Header.Set("X-TC-Region", provider.region)
	req.Header.Set("X-TC-Timestamp", strconv.FormatInt(time.Now().Unix(), 10))

	resp, err := provider.httpClient.Do(req)
	if err != nil {
		return nil, errors.Trace(err)
	}
	defer func() {
		_ = resp.Body.Close()
	}()
	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, errors.Errorf("get session token failed, status=%d body=%s", resp.StatusCode, string(respBody))
	}

	var parsedResp tencentCloudAssumeRoleResponse
	if err := json.Unmarshal(respBody, &parsedResp); err != nil {
		return nil, errors.Annotatef(err, "failed to unmarshal tencent cloud sts response: %s", string(respBody))
	}
	if parsedResp.Response.Error != nil {
		return nil, errors.Errorf("tencent cloud sts failed, code=%s message=%s request-id=%s",
			parsedResp.Response.Error.Code, parsedResp.Response.Error.Message, parsedResp.Response.RequestID)
	}
	if len(parsedResp.Response.Credentials.TmpSecretID) == 0 || len(parsedResp.Response.Credentials.TmpSecretKey) == 0 {
		return nil, errors.Errorf("tencent cloud sts returned empty credentials, request-id=%s", parsedResp.Response.RequestID)
	}
	expiredTime := parsedResp.Response.ExpiredTime
	if expiredTime == 0 {
		expiredTime = time.Now().Unix() + provider.durationSeconds
	}
	return &tencentCloudSessionCredentials{
		AccessKeyID:         parsedResp.Response.Credentials.TmpSecretID,
		SecretAccessKey:     parsedResp.Response.Credentials.TmpSecretKey,
		SessionToken:        parsedResp.Response.Credentials.Token,
		ExpirationTimestamp: expiredTime,
	}, nil
}