package engine

import (
	"encoding/base64"
	"encoding/json"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/aliyun/credentials-go/credentials"
)

// oidcCredentialWrap wraps aliyun OIDC credentials.Credential to fix AuthenticationFail.OIDCToken.Expired:
// credentials-go only refreshes based on Aliyun STS expiry; if STS session lifetime is longer than the
// Kubernetes service account JWT exp, AssumeRoleWithOIDC runs with an expired OIDC token.
// This wrapper busts the inner provider (new Credential = empty STS cache) when the JWT file claims
// are expired (file may have been rotated) or when STS returns OIDCToken.Expired.
type oidcCredentialWrap struct {
	mu        sync.Mutex
	inner     credentials.Credential
	tokenPath string
	factory   func() (credentials.Credential, error)
}

func newOIDCCredentialWrap(inner credentials.Credential, tokenPath string, factory func() (credentials.Credential, error)) credentials.Credential {
	return &oidcCredentialWrap{
		inner:     inner,
		tokenPath: tokenPath,
		factory:   factory,
	}
}

func (w *oidcCredentialWrap) GetAccessKeyId() (*string, error) {
	cm, err := w.getCredentialWithRefresh()
	if err != nil {
		return nil, err
	}
	return cm.AccessKeyId, nil
}

func (w *oidcCredentialWrap) GetAccessKeySecret() (*string, error) {
	cm, err := w.getCredentialWithRefresh()
	if err != nil {
		return nil, err
	}
	return cm.AccessKeySecret, nil
}

func (w *oidcCredentialWrap) GetSecurityToken() (*string, error) {
	cm, err := w.getCredentialWithRefresh()
	if err != nil {
		return nil, err
	}
	return cm.SecurityToken, nil
}

func (w *oidcCredentialWrap) GetBearerToken() *string {
	return w.inner.GetBearerToken()
}

func (w *oidcCredentialWrap) GetType() *string {
	return w.inner.GetType()
}

func (w *oidcCredentialWrap) GetCredential() (*credentials.CredentialModel, error) {
	return w.getCredentialWithRefresh()
}

func (w *oidcCredentialWrap) getCredentialWithRefresh() (*credentials.CredentialModel, error) {
	w.mu.Lock()
	defer w.mu.Unlock()

	if w.inner == nil {
		inner, err := w.factory()
		if err != nil {
			return nil, err
		}
		w.inner = inner
	}

	// If JWT claims say expired, kubelet may have already rotated the file — new Credential clears STS cache.
	if w.tokenPath != "" && jwtClaimsExpired(w.tokenPath) {
		inner, err := w.factory()
		if err != nil {
			return nil, err
		}
		w.inner = inner
	}

	cm, err := w.inner.GetCredential()
	if err != nil && isOIDCTokenExpiredErr(err) {
		inner, err2 := w.factory()
		if err2 != nil {
			return nil, err2
		}
		w.inner = inner
		return w.inner.GetCredential()
	}
	return cm, err
}

func isOIDCTokenExpiredErr(err error) bool {
	if err == nil {
		return false
	}
	msg := err.Error()
	return strings.Contains(msg, "OIDCToken.Expired") || strings.Contains(msg, "AuthenticationFail.OIDCToken.Expired")
}

// jwtClaimsExpired returns true if the file looks like a JWT and exp is before now (with small skew).
func jwtClaimsExpired(path string) bool {
	exp, ok := parseJWTExpiryUnix(path)
	if !ok {
		return false
	}
	// Allow 90s clock skew between kube issuer and this pod
	return time.Now().Unix() > exp+90
}

func parseJWTExpiryUnix(path string) (exp int64, ok bool) {
	b, err := os.ReadFile(path)
	if err != nil {
		return 0, false
	}
	s := strings.TrimSpace(string(b))
	parts := strings.Split(s, ".")
	if len(parts) != 3 {
		return 0, false
	}
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return 0, false
	}
	var claims struct {
		Exp float64 `json:"exp"`
	}
	if err := json.Unmarshal(payload, &claims); err != nil || claims.Exp == 0 {
		return 0, false
	}
	return int64(claims.Exp), true
}