package k8s

import (
	"context"
	"encoding/base64"
	"fmt"
	"strings"
	"time"

	corev1 "k8s.io/api/core/v1"
	apierr "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/aws-iam-authenticator/pkg/token"
	_ "sigs.k8s.io/aws-iam-authenticator/pkg/token"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	// tokenExpirationKey is used for retrieving token expiration timestamp from secret.
	tokenExpirationKey = "expirationTimestamp"
	// bearerTokenKey is used to retrieve bearer token from secret.
	bearerTokenKey = "token"
)

// EKSTokenManager manages the EKS apiserver bearer token.
type EKSTokenManager interface {
	RetrieveToken(ctx context.Context, opt EKSBearerTokenOpt) (string, error)
}

// eksTokenMngr manages the credentials for EKS cluster.
type eksTokenMngr struct {
	client.Client
	tokGen token.Generator
}

func NewEKSTokenManager(cli client.Client, tokGen token.Generator) EKSTokenManager {
	return &eksTokenMngr{Client: cli, tokGen: tokGen}
}

// EKSBearerTokenOpt specifies the options for retrieving bearer token for the
// EKS cluster.
type EKSBearerTokenOpt struct {
	EKSTarget
	// now is the oracle that returns current timestamp
	// it's not exposed and only for internal mocking/testing
	now func() time.Time
}

// RetrieveToken returns the temp bearer token for the given EKS cluster.
// If the cached token has expired, retrieve the latest token and update cache.
func (e *eksTokenMngr) RetrieveToken(ctx context.Context, opt EKSBearerTokenOpt) (string, error) {
	secret, err := e.retrieveCredSecret(ctx, opt)
	if err != nil && !apierr.IsNotFound(err) {
		return "", err
	}
	tokenStr, cacheHit := "", true
	if err != nil {
		secret, cacheHit = nil, false
	} else if expired, err := isSecretExpired(secret, opt.now); err != nil {
		return "", err
	} else if expired {
		cacheHit = false
	}

	if cacheHit {
		tokenStr, err = retrieveTokenString(secret)
	} else {
		g := e.tokGen
		if g == nil {
			g, err = token.NewGenerator(false, false)
			if err != nil {
				return "", err
			}
		}
		var tok token.Token
		tok, err = g.GetWithOptions(&token.GetTokenOptions{
			Region:               opt.Region,
			ClusterID:            opt.Cluster,
			AssumeRoleARN:        opt.AssumeRole,
			AssumeRoleExternalID: opt.ExternalID,
		})
		if err != nil {
			return "", err
		}
		tokenStr = tok.Token
		var expireStr string
		expireStr, err = metav1.NewTime(tok.Expiration).MarshalQueryParameter()
		if err != nil {
			return "", err
		}
		err = e.updateSecret(ctx, tok.Token, expireStr, secret, opt)
	}
	if err != nil {
		return "", err
	}
	return tokenStr, nil
}

func isSecretExpired(secret *corev1.Secret, now func() time.Time) (bool, error) {
	if now == nil {
		now = time.Now
	}
	expireStr, err := retrieveValueFromSecret(secret, tokenExpirationKey)
	if err != nil {
		return true, err
	}
	var expTime metav1.Time
	err = expTime.UnmarshalQueryParameter(expireStr)
	if err != nil {
		return true, err
	}
	return now().After(expTime.Time), nil
}

func retrieveTokenString(secret *corev1.Secret) (string, error) {
	return retrieveValueFromSecret(secret, bearerTokenKey)
}

func retrieveValueFromSecret(secret *corev1.Secret, key string) (string, error) {
	buf, err := base64.StdEncoding.DecodeString(string(secret.Data[key]))
	if err != nil {
		return "", err
	}
	return string(buf), nil
}

func base64Encode(str string) []byte {
	return []byte(base64.StdEncoding.EncodeToString([]byte(str)))
}

// updateSecret updates the token contents in the secret.
// If the secret is not created yet, create a new secret.
func (e *eksTokenMngr) updateSecret(ctx context.Context, tokenStr, expiration string, secret *corev1.Secret, opt EKSBearerTokenOpt) error {
	var secretNotExist bool
	if secret == nil {
		key, err := retrieveRoleCredSecretKey(opt)
		if err != nil {
			return err
		}
		secret, secretNotExist = &corev1.Secret{
			ObjectMeta: metav1.ObjectMeta{
				Name:      key.Name,
				Namespace: key.Namespace,
			},
		}, true
	}
	secret.Data = make(map[string][]byte)
	secret.Data[bearerTokenKey] = base64Encode(tokenStr)
	secret.Data[tokenExpirationKey] = base64Encode(expiration)
	if secretNotExist {
		return e.Create(ctx, secret)
	}
	return e.Update(ctx, secret)
}

// retrieveCredSecret returns the cred secret associated with the AWS role.
// The cred secret stores the bearer token and the expiration timestamp for
// accessing the EKS cluster. If the cred secret is nil, then use current AWS
// identity to access the EKS cluster.
func (e *eksTokenMngr) retrieveCredSecret(ctx context.Context, opt EKSBearerTokenOpt) (*corev1.Secret, error) {
	secret := &corev1.Secret{}
	key, err := retrieveRoleCredSecretKey(opt)
	if err != nil {
		return nil, err
	}
	err = e.Get(ctx, key, secret)
	return secret, err
}

// retrieveRoleCredSecretKey generates the secret identity based on role ARN
func retrieveRoleCredSecretKey(opt EKSBearerTokenOpt) (client.ObjectKey, error) {
	var account, role string
	_, err := fmt.Sscanf(
		strings.ReplaceAll(opt.AssumeRole, ":", " "),
		strings.ReplaceAll("arn:aws:iam::%s:role/%s", ":", " "),
		&account, &role)
	if err != nil {
		return client.ObjectKey{}, err
	}
	return client.ObjectKey{
		Name:      fmt.Sprintf("cred-%s-%s-%s-%s", account, role, opt.Region, opt.Cluster),
		Namespace: opt.Namespace,
	}, nil
}